Pytorch模型训练流程（简化标准版）

在 PyTorch 中训练一个模型通常遵循一个标准化的流程。我们将通过一个具体的例子——在 CIFAR10 数据集上训练一个简单的卷积神经网络（CNN）——来详细解释每一步。

PyTorch 训练模型的核心步骤概览

一个完整的训练流程可以分解为以下几个关键部分：

    数据准备 (Data): 加载并预处理数据，用Dataset 和 DataLoader进行封装。

    模型定义 (Model): 使用 torch.nn.Module 构建你的神经网络结构。

    损失函数定义 (Loss Function): 选择一个合适的损失函数来衡量模型输出与真实标签之间的差距。

    优化器定义 (Optimizer): 选择一个优化算法来根据损失函数的梯度更新模型的参数（权重）。

    训练循环 (Training Loop): 迭代数据，执行前向传播、计算损失、反向传播和参数更新。

    评估循环 (Evaluation Loop): 在验证/测试集上评估模型的性能，此过程不计算梯度。

    保存与加载模型 (Saving & Loading): 保存训练好的模型以备将来使用。

第1步：准备数据 (Data)

    数据是模型训练的基础。PyTorch 提供了Dataset和DataLoader这两个的工具来标准化数据加载流程。

    Dataset: 用于封装和访问数据集的样本和标签TorchVision 已经为我们内置了许多常用数据集（如CIFAR10,MNIST）。

    DataLoader: 它将Dataset包装起来，提供了批量（batching）、打乱（shuffling）和并行加载数据等功能，方便训练。

importtorch

importtorchvision

importtorchvision.transforms as transforms

 

# 1. 定义数据预处理方式

# transforms.Compose 可以将多个预处理步骤串联起来。

# CIFAR10 的图像是 3x32x32 的，我们将其标准化到 [-1, 1] 的范围。

transform =transforms.Compose(

    [transforms.ToTensor(), # 将 PIL Image 或 numpy.ndarray 转换为 FloatTensor，并缩放到 [0, 1]

     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]) # 标准化到 [-1, 1]

 

# 2. 下载并加载训练集和测试集

batch_size =64# 定义每个批次加载的图片数量

 

# 加载训练集

# download=True 会在 root 目录下自动下载数据集

trainset =torchvision.datasets.CIFAR10(root='./data', train=True,download=True, transform=transform)

 

trainloader =torch.utils.data.DataLoader(trainset, batch_size=batch_size,shuffle=True, num_workers=2) # shuffle=True 打乱数据

 

# 加载测试集

testset =torchvision.datasets.CIFAR10(root='./data', train=False,download=True, transform=transform)

 

testloader =torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False, num_workers=2)

 

# 类别标签

classes =('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

 

print("数据加载完成！")

第2步：定义模型 (Model)

    在 PyTorch 中，所有的神经网络模型都应该继承自torch.nn.Module。需要在 __init__ 方法中定义网络的层，在 forward 方法中定义数据的前向传播逻辑。

importtorch.nn as nn

importtorch.nn.functional as F

 

classSimpleCNN(nn.Module):

    def__init__(self):

        super(SimpleCNN, self).__init__()

        # 定义卷积层

        self.conv1 =nn.Conv2d(3, 16, 5) 

 # 输入通道3, 输出通道16, 卷积核大小5x5

        self.pool =nn.MaxPool2d(2, 2)  

 # 2x2 最大池化

        self.conv2 =nn.Conv2d(16, 32, 5) 

# 输入通道16, 输出通道32, 卷积核大小5x5



        # 定义全连接层

        self.fc1 =nn.Linear(32*5*5, 120) 

# 输入维度 32*5*5，输出120

        self.fc2 =nn.Linear(120, 84)

        self.fc3 =nn.Linear(84, 10) # 输出10，对应10个类别

 

    defforward(self, x):

        # 定义前向传播路径

        # x 的维度: [batch_size, 3, 32, 32]

        x =self.pool(F.relu(self.conv1(x))) 

# -> [batch_size, 16, 14, 14]

        x =self.pool(F.relu(self.conv2(x)))

# -> [batch_size, 32, 5, 5]

         

        # 将多维张量展平成一维向量

        x =torch.flatten(x, 1) # -> [batch_size, 32 * 5 * 5]

        x =F.relu(self.fc1(x)) # -> [batch_size, 120]

        x =F.relu(self.fc2(x)) # -> [batch_size, 84]

        x =self.fc3(x)        
#
  -> [batch_size, 10]

        returnx

 

# 实例化模型

net =SimpleCNN()

print("模型定义完成！")

第3步和第4步：定义损失函数和优化器

   损失函数: 对于多分类问题，最常用的是交叉熵损失nn.CrossEntropyLoss。注意，这个损失函数内部已经包含了Softmax 操作，所以我们模型的输出层不需要加Softmax。

   优化器: 随机梯度下降（SGD）。我们传入模型的参数net.parameters()和学习率lr来初始化它。

import torch.optim as optim

 
# 定义损失函数

criterion =nn.CrossEntropyLoss()

# 定义优化器

# 使用带动量的 SGD

optimizer=optim.SGD(net.parameters(),lr=0.001, momentum=0.9)

print("损失函数和优化器定义完成！")

第5步：训练循环 (Training Loop)

        这是整个流程的核心。我们会迭代（epoch）多次数据集，在每个epoch 中，将会分批次（batch）地处理所有数据。

训练循环中的关键步骤：

    1.optimizer.zero_grad(): 清空上一轮的梯度。PyTorch 的梯度是累加的，所以每轮训练前必须清零。

    2.outputs = net(inputs): 前向传播。将一个批次的数据输入模型，得到预测输出。

    3.loss = criterion(outputs, labels): 计算损失。比较模型输出和真实标签，计算损失值。

    4.loss.backward(): 反向传播。根据损失值，autograd 会自动计算模型中所有需要梯度的参数的梯度。

    5.optimizer.step(): 更新参数。优化器根据计算出的梯度来更新模型的权重。

# 决定使用 CPU 还是 GPU

device =torch.device("cuda:0"iftorch.cuda.is_available() else"cpu")

print(f"将使用 {device} 进行训练。")

 

# 将模型移动到指定设备

net.to(device)

 

# --- 开始训练循环 ---

num_epochs =5# 训练整个数据集的轮次

 

forepoch inrange(num_epochs):

     

    running_loss =0.0

    # trainloader 会返回批次的数据和标签

    fori, data inenumerate(trainloader, 0):

        # 1. 获取输入数据

        inputs, labels =data

         

        # 2. 将数据移动到指定设备

        inputs, labels =inputs.to(device), labels.to(device)

         

        # 3. 梯度清零

        optimizer.zero_grad()

         

        # 4. 前向传播

        outputs =net(inputs)

         

        # 5. 计算损失

        loss =criterion(outputs, labels)

         

        # 6. 反向传播

        loss.backward()

         

        # 7. 更新权重

        optimizer.step()

         

        # 打印统计信息

        running_loss +=loss.item()

        if(i +1) %500==0: # 每500个 mini-batches 打印一次

            print(f'[Epoch {epoch + 1}, Batch {i +
  1}] loss: {running_loss / 500:.3f}')

            running_loss =0.0

 

print('--- 训练完成 ---')

第6步：评估模型 (Evaluation)

训练完成后，我们需要在测试集上检验模型的泛化能力。

评估循环的关键点：

model.eval():将模型切换到评估模式。这会关闭 Dropout 和 BatchNorm 等在训练和测试时行为不同的层。

with torch.no_grad(): 在这个代码块中，所有计算都不会产生梯度。这可以节省内存并加速计算，因为我们只是在做前向传播，不需要反向传播。

correct =0

total =0

 

# 将模型切换到评估模式

net.eval() 

 

# 由于我们不需要计算梯度，使用 torch.no_grad()

with torch.no_grad():

    fordata intestloader:

        images, labels =data

        images, labels =images.to(device), labels.to(device)


        # 前向传播得到预测输出

        outputs =net(images)


        # 获取预测结果

        # torch.max 返回 (最大值, 最大值索引)

        _, predicted =torch.max(outputs.data, 1)

        total +=labels.size(0)

        correct +=(predicted ==labels).sum().item()

 

accuracy =100*correct /total

print(f'模型在测试集上的准确率: {accuracy:.2f} %')

第7步：保存和加载模型

通常我们只保存模型的状态字典（state dictionary），它包含了模型的所有参数（权重和偏置）。这种方式更轻量灵活。

# 保存模型

MODEL_PATH ='cifar_net.pth'

torch.save(net.state_dict(),
  MODEL_PATH)

print(f"模型已保存到 {MODEL_PATH}")

 

# --- 加载模型 ---

 

# 1. 重新实例化模型结构

loaded_net =SimpleCNN()

 

# 2. 加载状态字典

loaded_net.load_state_dict(torch.load(MODEL_PATH))

 

# 3. 将模型设置评估模式

loaded_net.to(device)

loaded_net.eval()

 

print("\n模型加载成功")

 

# 现在可以使用 loaded_net 进行预测了

以上就是训练的标准流程了，如果大家有不懂得可以评论区提出