pytorch编写训练、测试模型的一整个套路

以CIFAR10的网络模型为例。

网络结构

首先根据结构图编写网络结构：

import torch
from torch import nn

# 搭建神经网络
class Cifar10(nn.Module):
    def __init__(self):
        super(Cifar10, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(3, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 5, padding=2),
            nn.MaxPool2d(2),
            nn.Flatten(),
            nn.Linear(64*4*4, 64),
            nn.Linear(64, 10)
        )
    def forward(self, input):
        output = self.model(input)
        return output

# 验证网络的正确性
if __name__ == '__main__':
    cifar10 = Cifar10()
    input = torch.ones([64, 3, 32, 32])
    output = cifar10(input)
    print(output.shape)

单独运行此文件terminal打印：

tensor([64,10)]

64行数据代表64张（batch_size）图片，每行有10个数据代表每张图片在10个类别的概率分布。

训练文件

然后编写训练文件：

import torchvision
from torch.utils.tensorboard import SummaryWriter

from model import *   # 导入编写网络模型的文件
from torch.utils.data import DataLoader

# 定义训练的设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# 准备数据集
train_data = torchvision.datasets.CIFAR10(root="dataset", train=True, transform=torchvision.transforms.ToTensor(),
                                          download=True)
test_data = torchvision.datasets.CIFAR10(root="dataset", train=False, transform=torchvision.transforms.ToTensor(),
                                          download=True)

# 数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)

# 加载数据集
train_dataloader = DataLoader(dataset=train_data, batch_size=64)
test_dataloader = DataLoader(dataset=test_data, batch_size=64)

# 创建网络模型
cifar10 = Cifar10()
cifar10 = cifar10.to(device)

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

# 优化器
lr = 1e-2
optimizer = torch.optim.SGD(cifar10.parameters(), lr=lr)

writer = SummaryWriter("./logs")

epochs = 100
for epoch in range(epochs):
    print("----------开始第{}轮训练".format(epoch))
    # 开始训练步骤
    # 让网络进入训练状态 对于Dropout、BatchNormal等特殊层有作用，必须调用（没有的话也可以调用）
    cifar10.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = cifar10(imgs)
        loss = loss_fn(output, targets)

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

    # 开始验证
    # 让网络进入验证状态 对于Dropout、BatchNormal等特殊层有作用，必须调用（没有的话也可以调用）
    cifar10.eval()
    total_test_loss = 0
    preds = 0
    with torch.no_grad():
        for data in test_dataloader:
            imgs, targets = data
            imgs = imgs.to(device)
            targets = targets.to(device)
            output = cifar10(imgs)
            pred = (output.argmax(1) == targets).sum()
            preds += pred
            loss = loss_fn(output, targets)
            total_test_loss += loss.item()
    print("test loss:{}".format(total_test_loss))
    print("test accuracy:{}".format(preds/test_data_size))
    writer.add_scalar("test_loss", total_test_loss, epoch)
    writer.add_scalar("test_accuracy", preds/test_data_size, epoch)

    # 保存每轮模型训练结果
    torch.save(cifar10, "models/cifar10_{}.pth".format(epoch))

writer.close()

1.梯度清零：是为了上一轮计算的梯度不对本轮更新参数造成影响所以每轮要先进行梯度清零再根据梯度更新参数。
2.验证时不进行梯度计算：一是节省资源二是验证只是为了我们从各项指标参考模型训练的如何，得到的梯度不要对参数的更新造成影响。
3.预测准确率：
output.argmax(1)：是得到推理输出（预测值）每行数据最大值所在的位置，即得到预测的类别。

若参数为0则计算没列数据最大值所在位置，但这里我们一行为一张图片的概率分布，所以我们要横向计算而不是纵向的。

output.argmax(1) == targets：预测的类别与真实的类别进行对比，同一张照片的预测类和真实类相等则为True。

(output.argmax(1) == targets).sum()：将每个预测相等的值进行相加，再处以测试数据的总数，即得到在测试数据上的预测准确率。

运行结果

在terminal输入启动tensorboard的命令：tensorboard --logdir=logs，然后打开对应网址：
首先看在验证集上的准确率：

总体上升算是正常，但是到0.65就不再上升还算是挺低的（但也只是个实验非实际应用）。
再看train_loss：

loss下降也还算正常。
再看test_loss：

这个曲线走势就不正常了，20轮以后开始飙升。结合train_loss走势来看，可能是局部最优了，网络太简单学习率后期没下降等导致过拟合，也就是网络钻牛角尖了往局部方向拼命学习。

修改

这里浅浅做了个实验（只调学习率），看是否会有改善。
将优化器的使用语句放入100轮的循环当中，然后每轮学习率×0.95进行下降。

先看正确率：

训练100轮到了最后仍然还有上升的趋势，而且接近于0.7比刚才更高，说明可以看到效果了。
再看train_loss：

曲线非常的平滑，对比刚才最后稍微的震荡是更好的。
最后来看刚才出问题的test_loss：

非常的漂亮，也是平滑的下降。
至此可以看到我们的改动还是很成功的。

测试文件

首先随便在网上找了张cifar10数据集的10个类别当中的类的图片，这里我找了张狗举例。

import torch
from PIL import Image
import torchvision

img_path = "./dog.png"
img = Image.open(img_path)

img = img.convert("RGB")  # 保留rgb通道 png有4个通道，多一个透明度通道 这样就可以适应各种格式的图片

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
                                            torchvision.transforms.ToTensor()])

img = transform(img)
print(img.shape)

img = torch.reshape(img, (1, 3, 32, 32))

img = img.cuda()
model = torch.load("models/cifar10_99.pth")  # 因为观察曲线走势发现最后一轮的效果最好所以用最后一次训练保存的模型
model.eval()
with torch.no_grad():     # 节约内存
    output = model(img)
print(output.argmax(1))

问题

这里会有几个问题：
1.Resize方法里面要是一个参数，所以用元组或者list括起来。因为一个int表示最小边resize成指定值，最大边等比例缩放；如果是给两个边大小，即缩放成指定的[h, w]。
2.因为这里输入的只有一张图片，不是文件夹内的批量图片，所以没用dalaloader设置batch_size的批量获取。因此如果不用reshape成四维（多一个batch）的话会报这个错：

RuntimeError: Expected 4-dimensional input for 4-dimensional weight [32, 3, 5, 5],
 but got 3-dimensional input of size [3, 32, 32] instead

也就是网络输入的应该是4维，但实际只有3维。
3.如果输入的数据没有使用cuda的话，会报一个这样的错：

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) 
should be the same or input should be a MKLDNN tensor and weight is a dense tensor

输入类型与输出类型应当相同，因为权重我们是使用cuda进行训练，所以这里我们也需要cuda类型的输入。
如果是在不同的环境上跑模型的话，比如有gpu的设备上训练，到只有cpu的设备上测试，加载模型的时候就需要指定映射环境，映射到cpu的环境中：

model = torch.load("models/cifar10_99.pth", map_location=torch.device("cpu"))

输出

输出结果（cpu输出）

tensor([5])

gpu输出

tensor([5], device='cuda:0')

到pytorch官网看一下cifar10数据集对应的类：

可以看到idx为5的类为dog，预测正确。

总结