【PyTorch】构造VGG19网络进行本地图片分类（超详细过程）——项目介绍

本篇博客主要解决以下3个问题：

1. 如何自定义网络（以VGG19为例）。
2. 如何自建数据集并加载至模型中。
3. 如何使用自定义数据训练自定义模型。

第一篇：【PyTorch】构造VGG19网络进行本地图片分类（超详细过程）——项目介绍
第二篇：【PyTorch】构造VGG19网络进行本地图片分类（超详细过程）——程序代码
Github：https://github.com/MarvelInSky/vgg_classify

---

---

一、VGG简介

VGG的网络结构如下，本篇博客以VGG19（E列），过多内容不再介绍。

二、数据集介绍

Animal Image Dataset(DOG, CAT and PANDA)
Dataset for Image Classification Practice

下载地址：https://www.kaggle.com/ashishsaxena2209/animal-image-datasetdog-cat-and-panda

该数据集共包含3类目标：狗、猫和熊猫；每种图片各1000张；图片尺寸不固定；大部分图片为RGB图片，少部分图片为灰度图片，所以在处理数据的时候要注意通道数。

将其分为train、test两个文件：

• train：用于训练模型，存放每种类别的第1至第900张图片，共包含2700张。
• test：用于验证和测试，存放每种类别的第901至第1000张图片，共包含300张。

三、项目程序

3.1 测试环境

测试环境

Python3.8
Cuda10.1
PyTorch1.7
1
2
3

所需依赖

matplotlib
torch
torchvision
1
2
3

3.2 程序结构

3.3 常用函数

3.3.1 Module类

自定义一个模型——通过继承nn.Module类来实现，需要在__init__构造函数中申明各层的定义，在forward中实现层之间的连接关系，实际上就是前向传播的过程。

import torch.nn as nn

class VGG(nn.Module):

    # 初始化并定义网络结构
    def __init__():
    	# 定义网络结构
    
    def forward(self, x)
    	# 前项传播的方法
1
2
3
4
5
6
7
8
9
10

参考：
https://blog.csdn.net/qq_27825451/article/details/90705328

3.3.2 Dataset类

pytorch包里提供了Dataset来对数据进行加载，通过继承Dataset类得到MyDataset类，我们可以实现对自己数据集的加载。继承的类需要重写__init__、__len__、__getitem__。

__len__：返回数据内包含的数据总量，可以将所有数据的地址存储在一个列表里，返回字典的长度即刻。
__getitem__：根据对象的索引，返回(x,y)，x为图像的数据，y为标签。
__init__：根据以上__len__和__getitem__的需求，可以将每张数据图片的路径存在一个list中，方便getitem的调用。

参考：
https://blog.csdn.net/weixin_44168899/article/details/100929727
https://blog.csdn.net/leviopku/article/details/99958182

3.3.3 DataLoader类

在上一步中，通过继承Dataset类得到MyDataset获得读取数据的方法，现在使用DataLoader加载这些类。

train_loader = DataLoader(dataset, batch_size, num_workers=2, shuffle=True)
1

dataset：为实例化MyDataset得到的对象。
batchsize：就是每次从数据中心加载进模型的数量。
num_workers：相当于有两个线程在同时处理加载。
shuffle=True：表示加载数据的时候是乱序。

参考：
https://blog.csdn.net/zw__chen/article/details/82806900
https://pytorch.org/docs/1.1.0/_modules/torch/utils/data/dataloader.html

3.3.4 损失函数

定义：损失函数使用交叉熵

loss_function = torch.nn.CrossEntropyLoss()
1

在每一batch中使用，计算误差并进行反向传播

loss = loss_function(outputs, labels)  # outputs为模型得到的值，labels为真值
loss.backward()  # 将loss进行反向传播
1
2

表示使用交叉熵作为损失函数。

3.3.5 优化器

定义：使用SGD优化器，scheduler自定义调整学习率

optimizer = torch.optim.SGD(net.parameters(), lr, momentum=0.9, weight_decay=5e-4)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones, gamma=0.1)
1
2

参考：https://blog.csdn.net/zisuina_2/article/details/103258573

3.3.6 optimizer.zero_grad()

optimizer.zero_grad()意思是把梯度置零，即将loss关于weight的导数变成0.

参考：https://blog.csdn.net/scut_salmon/article/details/82414730

3.3.7 net.train()与net.eval()

net.train()：启用Batch Normalization和Dropout
net.eval()：不启用Batch Normalization和Dropout

参考：https://blog.csdn.net/qq_38410428/article/details/101102075

3.3.8 保存与加载模型

# 加载自定义模型
net = vgg_model.VGG(img_siz, input_channel, num_class)

# 训练
...

# 训练完成后保存模型
torch.save(net.state_dict(), weights_path)
1
2
3
4
5
6
7
8

net.state_dict()表示只保存模型中的参数，可以减小模型的体积。

# 加载自定义模型，此时里面参数是随机的
net = vgg_model.VGG(img_siz, input_channel, num_class)

# 加载训练好的参数到模型
net.load_state_dict(torch.load(opt.model_path))
1
2
3
4
5

3.3.9 tensorboard

tensorboard可用于记录训练时的日志，可以在训练时，实时观看相关数据的变化曲线，了解模型的变化趋势。

from torch.utils.tensorboard import SummaryWriter

# 开启记录
writer = SummaryWriter(log_dir)

# 添加记录数据
for epoch in range(1,epochs)
	...
	writer.add_scalar('Average loss', loss, epoch)
    writer.add_scalar('Accuracy', acc, epoch)
1
2
3
4
5
6
7
8
9
10

使用方法：

writer.add_scalar('曲线图名称', 纵坐标, 横坐标)
1

查看方法：

在日志目录下打开终端输入命令：tensorboard --logdir=日志所在目录

启动后，在浏览器打开http://localhost:6006/ 。

3.3.10 tqdm动态进度条

希望在每一batch都显示实时当前的epoch、loss和lr，但是又不希望不停的print的这些值，因为这些print会减慢训练速度，也不方便查看其它epoch的效果。所以使用tqdm库的创建动态的进度条。

# 获取数据的总量
nb = len(train_dataset)

# 创建进度条 
# train_loader 需要遍历的数据
# total=数据的总量，为总量/batchsize
pbar = tqdm(enumerate(train_loader), total=int(nb / opt.batch_size))

# 在训练时
for step, (images, labels) in pbar:
	...
1
2
3
4
5
6
7
8
9
10
11

效果如下：

四、Debug

4.1 参数类型和模型权重类型不一致

RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
1

原因：在编写项目的过程中，考虑了是否使用GPU的情况；所以，在传递传递参数时，如果使用GPU，则x=x.cuda()，然而在初始化模型时，没有使用cuda()；这就导致参数的格式和初始化得到的模型中的权重格式不一致。

解决：

net = vgg_model.VGG(input_channel=3, num_class=opt.num_class)

# 改为
if opt.gpu:
    net = vgg_model.VGG(input_channel=3, num_class=opt.num_class).cuda()
else:
    net = vgg_model.VGG(input_channel=3, num_class=opt.num_class)
1
2
3
4
5
6
7

4.2 前项传播时卷积层内数据维度不匹配

# 在卷积层处发生错误
RuntimeError: Given groups=1, weight of size [512, 512, 3, 3], expected input[16, 256, 28, 28] to have 512 channels, but got 256 channels instead
1
2

原因：输入的尺寸和某一层的结构不匹配。

解决：在前项传播中，输出每一层的结构和每一层的输出结果，发现到第九层时出现错误，仔细观察每一层的input_channel和output_channel，发现第9成的网络输入输出为（512,512），此处由于疏忽写错了，将其改为（256, 512）即可。

4.3 前项传播时全连接层数据维度不匹配

# 在全连接层处发生错误
RuntimeError: mat1 dim 1 must match mat2 dim 0
1
2

原因：在参考代码中，第一层全连接的代码：nn.Linear(512, 4096)是将512个数据转为4096个数据，采用的数据集为cifar100，其像素尺寸为32X32，经过16层卷积和5次池化，可将原尺寸[3,32,32]转变为[512,1,1]，其中512X1X1=512。然而，我的数据集的像素尺寸已经resize成224X224，经过16个层卷积和5次池化，原尺寸[3,512,512]转变为[512,7,7]，其中51277=25088，所以第一层全连接的代码为：nn.Linear(25088, 4096)。

解决：但是这样写兼容性就差了许多，不能兼容各像素的图片，所以将25088这个值改成由输入图片像素构成的：$convOut=512\times (imgSize/24{)}^2$

4.4 GPU显存不足

cuda out of memory
1

原因：电脑的显存不够。

解决：原本输入图像的尺寸为224224，现改为128128，原本的batchsize为8改为4。

4.5 RGB图像中混入灰度图片

RuntimeError: Given groups=1, weight of size [64, 3, 3, 3], expected input[1, 1, 64, 64] to have 3 channels, but got 1 channels instead
1

原因：默认数据集中均位3通道图片（RGB图像），但其中有混入的单通道图片（灰度图片）。

解决：在使用Dataset加载数据时，判断是否为单通道，若为单通道图片则转为三通道图片。

4.6 训练过程中模型发散

Training Epoch: 1 [2/2700]	Loss: 1.2244	LR: 0.100000
Training Epoch: 1 [4/2700]	Loss: 20.7208	LR: 0.100000
Training Epoch: 1 [6/2700]	Loss: 25.6863	LR: 0.100000
Training Epoch: 1 [8/2700]	Loss: 0.0000	LR: 0.100000
Training Epoch: 1 [10/2700]	Loss: 1070.9552	LR: 0.100000
Training Epoch: 1 [12/2700]	Loss: 91453.2031	LR: 0.100000
Training Epoch: 1 [14/2700]	Loss: 489251648.0000	LR: 0.100000
Training Epoch: 1 [16/2700]	Loss: 372763991191060480.0000	LR: 0.100000
Training Epoch: 1 [18/2700]	Loss: nan	LR: 0.100000
Training Epoch: 1 [20/2700]	Loss: nan	LR: 0.100000


之后Loss均为nan
1
2
3
4
5
6
7
8
9
10
11
12
13

原因：在训练的时候很快就开始发散了，模型未收敛。

解决：将学习率（Learning rate，lr）从0.1改为0.0001。

参考博客

1. pytorch实现vgg19 训练自定义分类图片：https://www.cnblogs.com/wuzaipei/p/12652932.html
2. pytorch-cifar100： https://github.com/weiaicunzai/pytorch-cifar100
3. yolov5：https://github.com/ultralytics/yolov5