自己训练一个模型原来如此简单，AI时代知识获取太容易了

我爱学大模型

墨衍会员 · AI 创作全网分发码龄3年

关注

原创 2025-03-07 17:58:31 1.8k 阅读 AI 写同款· GEO 优化 ›

Ollama 可以在 Windows 上运行了阅读详情

最近在看pytorch的东西，于是想问一下deepseek，让其描述完整过程，果然不失所望。

Q：怎么用pytorch训练一个模型，并跑起来。

A：（下面的文章是ds回答，我用的元宝满血版）

使用PyTorch训练并运行一个模型的完整流程可分为以下步骤，结合多篇实践指南整理而成：

一、环境搭建

安装PyTorch

pip install torch torchvision torchaudio  # CPU版本
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117  # GPU版本（需CUDA）

验证安装：

import torch
print(torch.__version__, torch.cuda.is_available())
# 输出版本及GPU是否可用

虚拟环境（可选）推荐使用Miniconda管理环境：

conda create -n pytorch_env python=3.10
conda activate pytorch_env

二、数据准备

数据集加载使用内置数据集（如MNIST）或自定义数据集：

from torchvision import datasets, transforms

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 图像归一化
])
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_data = datasets.MNIST(root='./data', train=False, transform=transform)

数据分批（DataLoader）

from torch.utils.data import DataLoader
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)

三、模型定义

继承nn.Module构建自定义模型：

import torch.nn as nn

class SimpleNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28*28, 128)  # 输入层到隐藏层
        self.fc2 = nn.Linear(128, 10)     # 隐藏层到输出层

    def forward(self, x):
        x = x.view(-1, 28*28)  # 展平图像
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleNN()

四、损失函数与优化器

criterion = nn.CrossEntropyLoss()  # 分类任务常用交叉熵损失
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

五、训练循环

device = torch.device("cuda" if torch.cuda.is_available() else"cpu")
model.to(device)  # 将模型移至GPU（若有）

for epoch in range(10):
    model.train()
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()        # 清空梯度
        outputs = model(inputs)      # 前向传播
        loss = criterion(outputs, labels)
        loss.backward()              # 反向传播
        optimizer.step()            # 更新参数

    # 每个epoch结束后评估
    model.eval()
    with torch.no_grad():
        correct = 0
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            correct += (predicted == labels).sum().item()
        accuracy = correct / len(test_data)
        print(f"Epoch {epoch+1}, Test Accuracy: {accuracy:.2f}")

六、模型保存与加载

保存模型权重

torch.save(model.state_dict(), 'model.pth')

加载模型

model = SimpleNN()
model.load_state_dict(torch.load('model.pth'))
model.eval()  # 切换为评估模式

七、高级技巧

GPU加速将数据和模型通过.to(device)移至GPU，显著提升训练速度。
防止过拟合

#添加Dropout层：
nn.Dropout(0.5)

# 2正则化：
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)

学习率调整

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
scheduler.step()  # 每个epoch后调用

可视化训练过程使用TensorBoard跟踪损失和准确率。

八、完整示例（线性回归）

# 生成数据
X = torch.randn(100, 1)
y = 3 * X + 2 + torch.randn(100, 1)

# 定义模型
class LinearRegression(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)
    def forward(self, x):
        return self.linear(x)

model = LinearRegression()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 训练
for epoch in range(100):
    outputs = model(X)
    loss = criterion(outputs, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (epoch+1) % 10 == 0:
        print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}')