写在前面
前段时间在 Git 上下载了 yolov5 的代码，经过调试，最后运行成功。但是发现对网络训练的步骤其实很不熟悉，于是乎最近看了看基于 pytorch 的深度学习——通过学习，对 pytorch 的框架有了较清晰的认识，也可以自己来构建一些模型来进行训练。如果你也发现自己只知道在 Git 上克隆别人的代码，但是自己对程序的结构不了解，那么下面的内容可能会帮到你！！！
这部分内容主要是根据B站视频总结而来，视频中给出了 pytorch 从安装到最后训练模型的完整教程，本篇文章主要总结神经网络的完整的模型训练套路，希望通过本篇文章可以让你对网络训练步骤有一个清晰的认识。
本次内容用到的数据集是CIFAR10，使用这个数据的原因是这个数据比较轻量，基本上所有的电脑都可以跑。CIFAR10 数据集里是一些 32X32 大小的图片，这些图片都有一个自己所属的类别（如 airplane、cat 等），如下图所示：
注意：这个数据集不需另外要从网页下载，程序中可以调整代码参数进行下载
我们先来了解一下我们需要进行的工作及实现的功能：我们首先需要下载数据集，然后通过数据来训练模型，并在测试集上进行测试，这时候我们可以保存我们训练好的模型。最后通过我们训练的模型来判断一些图片的类别（从网络上下载一些图片，判断它是猫是狗或是其他的类型【当然这个数据集只有 10 种类型，如上图所示的 10 种】）
下面我们就来一步步的介绍！！！【代码我分流程分部分介绍，完整代码放在文末自取】
完整网络模型训练步骤
1、准备数据集
很显然，没有数据一切都是空谈，那么第一步就是准备我们需要的数据集 CIFAR10。
#1、准备数据集train_dataset = torchvision.datasets.CIFAR10("./data", train=True, transform=torchvision.transforms.ToTensor(), download=Ture)test_dataset = torchvision.datasets.CIFAR10("./data", train=False, transform=torchvision.transforms.ToTensor(), download=Ture)
复制代码
第一个参数“./data”是指定下载数据集保存的位置，第二个参数 train=True/Flase 是指下载的数据是训练集数据还是测试集数据【True 表示训练集，Flase 表示测试集】，第三个参数是图片的一个转化，要将图片格式转化为 tensor 类型，第四个参数 download 为 True 表示你没有这个数据，这时候会自动下载数据，为 Flase 表示有这个数据，不会再进行下载【注意：这个参数设置成 True 且你有数据集，那同样不会进行数据下载，故这个参数一直设置成 True 就好了】。
我们可以打印数据集的长度来看一下这个数据集的大小，可以发现训练集有 5000 张图片，测试集有 1000 张图片。
train_dataset_size = len(train_dataset)test_dataset_size = len(test_dataset)print("train_dataset_size:{}".format(train_dataset_size))print("test_dataset_size:{}".format(test_dataset_size))
复制代码
2、加载数据集
#2、加载数据集train_dataset_loader = DataLoader(dataset=train_dataset, batch_size=64)test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=64)
复制代码
在得到数据集后，我们还要对数据集进行加载，加载数据集就类似于打包，比如这里的第二个参数设置的是 batch_size=64,则表示把 dataset 中的 64 个数据打包一起放入 dataloader 中。
3、搭建神经网络✨✨✨
加载好数据后，就可以搭建神经网络了，我们可以百度 CIFAR10 model,可以出现很多 CIFAR10 的网络模型，如图所示：
我们可以根据上图来搭建网络模型，如下：
#3、搭建神经网络class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.model1 = 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(1024, 64), nn.Linear(64, 10)
)
def forward(self, input): input = self.model1(input) return input
复制代码
这部分代码完全是根据上图中的模型一步步写的，具有一一对应的关系，只是在卷积中的 padding 需要我们根据前后输入输出的尺寸进行计算，最后发现三步卷积 padding 都为 2，这里给出 pytorch 官网的相关计算公式：
4、创建网络模型
这步只要一行代码，其实就是实列化了一个对象。
我们可以打印出来看一看我们自己创建的网络模型，如下图。可以看出和前文的结构是一致的。
到这里我们已经创建好了自己的模型，这个模型输入是 3x32x32 的图片【可以认为就是一个 3x32x32 的张量】，输出是 1x10 的向量。每当我们创建好一个模型后，应该检测一下模型的输入输出是否是我们所期待的，若不是则即使调整模型。我们可以用以下代码来检测输出是否符合要求。
net = Net()input = torch.ones((64, 3, 32, 32)) #64为batch_size,3x32x32表示张量尺寸output = net(input)print(output.shape)
复制代码
可以看出输出是符合要求的，64 是输入的 batch_size,相当于输入 64 张图片。
5、设置损失函数、优化器
设置损失函数、优化器这些都是神经网络的一些基础知识，不知道的自行补充。当然这里的损失函数和优化器可以和我不同，感兴趣的也可以改变这些来看看我们最后训练的效果会不会发生变化【我测试了几个，对于本例效果差别不大】
#5、设置损失函数、优化器#损失函数loss_fun = nn.CrossEntropyLoss() #交叉熵loss_fun = loss_fun.to(device)
#优化器learning_rate = 1e-2optimizer = torch.optim.SGD(net.parameters(), learning_rate) #SGD:梯度下降算法
复制代码
6、设置网络训练中的一些参数
这部分主要是用来记录一些训练测试的次数及网络训练轮数。
#6、设置网络训练中的一些参数total_train_step = 0 #记录总计训练次数total_test_step = 0 #记录总计测试次数epoch = 10 #设计训练轮数
复制代码
7、开始训练网络✨✨✨
进行网络训练时，我们首先会通过自己构建的网络得到输出，然后比较输出和真实值，计算出损失，最后通过反向传播，调整网络中参数的值。对于反向传播不理解的可以参考我的这篇文章：BP神经网络
#7、开始进行训练for i in range(epoch): print("---第{}轮训练开始---".format(i+1))
net.train() #开始训练，不是必须的，在网络中有BN，dropout时需要 for data in train_dataset_loader: imgs, targets = data targets = targets.to(device) outputs = net(imgs)
#比较输出与真实值，计算Loss loss = loss_fun(outputs, targets)
#反向传播，调整参数 optimizer.zero_grad() #每次让梯度重置 loss.backward() optimizer.step()
total_train_step += 1
if total_train_step % 100 == 0: print("---第{}次训练结束, Loss:{})".format(total_train_step, loss.item()))
复制代码
8、开始测试网络✨✨✨
对网络进行测试过程和训练是类似的，不同的是测试过程不需要通过反向传播来更新参数。
#8、开始进行测试,测试不需要进行反向传播net.eval() #开始测试，不是必须的，在网络中有BN，dropout时需要with torch.no_grad(): #这句表示测试不需要进行反向传播，即不需要梯度变化【可以不加】 total_test_loss = 0 #测试损失 total_test_accuracy = 0 #测试集准确率 for data in test_dataset_loader: imgs, targets = data outputs = net(imgs)
#计算测试损失 loss = loss_fun(outputs, targets) total_test_loss = total_test_loss + loss.item() accuracy = (outputs.argmax(1) == targets).sum() total_test_accuracy = total_test_accuracy + accuracyprint("第{}轮测试的总损失为：{}".format(i+1, total_test_loss))print("第{}轮测试的准确率为：{}".format(i+1, total_test_accuracy/test_dataset_size))
复制代码
9、保存模型
将每一个 epoch 的模型都保存下来，为后面物体识别准备模型。
#9、保存模型torch.save(net, "./self_model_{}".pth.format(i+1))print("模型已保存")
复制代码
检测训练模型的效果
介绍到这里，完整的自建网络模型训练步骤我们就讲完了，接下来来看看我们用之前保存的模型来检测一些我们从网络上下载的图片，代码如下：
import torchimport torchvisionfrom PIL import Imagefrom torch import nn
image_path = "./imgs/airplane.png" #网络下载的图片放置地址image = Image.open(image_path)
image = image.convert('RGB') #将图片转化为RGB三通道图片，有的图片有4个通道（多了个透明度）
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)), torchvision.transforms.ToTensor()])
image = transform(image)
class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.model1 = 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(1024, 64), nn.Linear(64, 10)
)
def forward(self, x): x = self.model1(x) return x
model = torch.load("net_29.pth", map_location=torch.device('cpu'))print(model)
image = torch.reshape(image, (1, 3, 32, 32))model.eval()with torch.no_grad(): output = model(image)
print(output.argmax(1))
复制代码
网络下载图片如下：
输出结果如下：
0 表示的就是 airplane【可以从官网中 10 种类型顺序得出，从上到下是 0-9】。
我们可以在来测试一张狗的图片，从官网可知，输出 5 为狗，原始图片和输出图片如下：
这里我们可以来看一下模型的检测损失和正确率（设置的 epoch=20），准确率大概在 65%左右。【这里是在Google Colab上用 GPU 训练的，单用 CPU 训练速度还是很慢】
模型的准确率似乎就停留在 65%上下，我尝试增大 epoch 到 30，但是准确率基本一致。同时我也用 3x3 的小卷积核代替 5x5 的卷积核、用卷积代替池化，用卷积代替全连接层等方式进行训练，但是效果都不显著，当然这里我只训练了 30 个 epoch,增大 epoch 效果可能会好，但耗时会比较多，这部分主要是学习训练模型的思路，感兴趣可以尝试各种方式看能否改进模型效果。
下图是用 Tensorboard 画的损失和准确率的曲线图，上文的代码中只关注模型的训练步骤，没有设计 tensorboard 的讲解，在文末源代码中会包含这部分内容。
完整代码
咻咻咻咻~~duang~~点个赞呗
评论