PyTorch学习笔记（8）训练模型全过程

终于到了训练的时候了，从今开始正式炼丹。
这里以 CIFAR10 数据集（常见的分类问题，共10个类别）为例，记录从准备数据集、搭建模型、训练模型的全过程，并做可视化展示。

---

准备数据集

train=True表示下载训练集，使用 transforms.ToTensor() 进行转化是因为计算机只能处理数字，图像也得转化为数字才能进行特征提取与后续复杂操作。
batch_size=64 可以理解为一次处理 64 张图片，如果 GPU 专用内存不足，需要调小。

import torchvision
from torch.utils.data import DataLoader

# 准备数据集
train_data = torchvision.datasets.CIFAR10("../dataset_CIFAR10",train=True,
                                          transform=torchvision.transforms.ToTensor(),download=True)
test_data  = torchvision.datasets.CIFAR10("../dataset_CIFAR10",train=False,
                                          transform=torchvision.transforms.ToTensor(),download=True)
# 打印数据长度
train_data_length = len(train_data)
test_data_length = len(test_data)
print("训练数据集的长度:{}".format(train_data_length))
print("测试数据集的长度:{}".format(test_data_length))

# 利用DataLoader来加载数据
train_dataloader = DataLoader(train_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)

准备网络模型

这里直接用之前复现的网络：PyTorch学习笔记（5）搭建简单神经网络

from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, Flatten

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.moudle = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x =  self.moudle(x)
        return x

# 搭建神经网络
tudui = Tudui()

定义损失函数与优化器

optim.SGD 表示使用随机梯度下降算法。

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 0.01 #学习速率
optimizer = torch.optim.SGD(params=tudui.parameters(),lr=learning_rate)

迭代训练过程

for i in range(epoch):
    print("---第{}轮训练开始---".format(i+1))

    #训练步骤如下
    for data in train_dataloader:
        imgs,targets = data
        outputs = tudui(imgs)
        loss = loss_fn(outputs,targets)     #两个参数的顺序不能错

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        # 逢100打印一次
        if total_train_step % 100 == 0:
            print("训练次数为：{},loss为：{}".format(total_train_step,loss))

训练过程如下：

再加上测试代码：

 # 测试步骤
    total_test_loss = 0
    total_accuracy = 0      #总共正确的个数
    #没有梯度，不进行优化，使用上一刻训练的模型
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets = data
            outputs = tudui(imgs)
            loss = loss_fn(outputs,targets)

            total_test_loss += loss
            total_test_step += 1
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
            pass
    print("本轮测试集上的Loss为:        {}".format(total_test_loss))
    print("本轮测试集上的正确率为:       {}".format(total_accuracy / test_data_length))
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    writer.add_scalar("test_accuracy",total_accuracy/test_data_length,total_test_step)

完整代码

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, Flatten


# 准备数据集
train_data = torchvision.datasets.CIFAR10("../dataset_CIFAR10",train=True,
                                          transform=torchvision.transforms.ToTensor(),download=False)
test_data  = torchvision.datasets.CIFAR10("../dataset_CIFAR10",train=False,
                                          transform=torchvision.transforms.ToTensor(),download=False)
# 打印数据长度
train_data_length = len(train_data)
test_data_length = len(test_data)
print("训练数据集的长度:{}".format(train_data_length))
print("测试数据集的长度:{}".format(test_data_length))

# 利用DataLoader来加载数据
train_dataloader = DataLoader(train_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)

# 搭建神经网络
class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.moudle = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x =  self.moudle(x)
        return x
    
# 搭建神经网络
tudui = Tudui()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 0.01 #学习速率
optimizer = torch.optim.SGD(params=tudui.parameters(),lr=learning_rate)

# 记录训练网络的一些参数
# 训练和测试的次数
total_train_step = 0
total_test_step = 0
# 训练的轮数
epoch = 10
# 添加tensorboard
writer = SummaryWriter("../logs")

for i in range(epoch):
    print("---第{}轮训练开始---".format(i+1))

    #训练步骤如下
    for data in train_dataloader:
        imgs,targets = data
        outputs = tudui(imgs)
        loss = loss_fn(outputs,targets)     #两个参数的顺序不能错

        # 优化器优化模型
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_train_step += 1
        # 逢100打印一次
        if total_train_step % 100 == 0:
            print("训练次数为：{},loss为：{}".format(total_train_step,loss))
            writer.add_scalar("train_loss",loss,total_train_step)   # y轴和x轴

    # 测试步骤
    total_test_loss = 0
    total_accuracy = 0      #总共正确的个数
    #没有梯度，不进行优化，使用上一刻训练的模型
    with torch.no_grad():
        for data in test_dataloader:
            imgs,targets = data
            outputs = tudui(imgs)
            loss = loss_fn(outputs,targets)

            total_test_loss += loss
            total_test_step += 1
            accuracy = (outputs.argmax(1) == targets).sum()
            total_accuracy += accuracy
            pass
    print("本轮测试集上的Loss为:        {}".format(total_test_loss))
    print("本轮测试集上的正确率为:       {}".format(total_accuracy / test_data_length))
    writer.add_scalar("test_loss",total_test_loss,total_test_step)
    writer.add_scalar("test_accuracy",total_accuracy/test_data_length,total_test_step)

    torch.save(tudui,"moudel_{}.pth".format(i+1))

writer.close()

查看结果

可以看到，损失 loss 在不断地下降。

使用 tensorboard 查看可以更加清晰地观测训练过程（也方便做 PPT）。