pytorch图像分类篇：2.pytorch官方demo实现一个分类器(LeNet)

参考内容来自：

• up主的b站链接：霹雳吧啦Wz的个人空间-霹雳吧啦Wz个人主页-哔哩哔哩视频
• up主将代码和ppt都放在了github：GitHub - WZMIAOMIAO/deep-learning-for-image-processing: deep learning for image processing including classification and object-detection etc.
• up主的CSDN博客：https://blog.csdn.net/qq_37541097/article/details/103482003

---

pytorch官网入门demo——实现一个图像分类器

demo的流程:

1.model.py——定义LeNet网络模型

2.train.py——加载数据集并训练，训练集计算loss,测试集计算accuracy,保存训练好的网络参数

3.predict.py——得到训练好的网络参数后，用自己找的图像进行分类测试

---

1.model.py

先给出代码，模型是基于LeNet做简单修改，层数很浅，容易理解：

# 使用torch.nn包来构建神经网络.
import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module): 					# 继承于nn.Module这个父类
    def __init__(self):						# 初始化网络结构
        super(LeNet, self).__init__()    	# 多继承需用到super函数
        self.conv1 = nn.Conv2d(3, 16, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):			 # 正向传播过程
        x = F.relu(self.conv1(x))    # input(3, 32, 32) output(16, 28, 28)
        x = self.pool1(x)            # output(16, 14, 14)
        x = F.relu(self.conv2(x))    # output(32, 10, 10)
        x = self.pool2(x)            # output(32, 5, 5)
        x = x.view(-1, 32*5*5)       # output(32*5*5)
        x = F.relu(self.fc1(x))      # output(120)
        x = F.relu(self.fc2(x))      # output(84)
        x = self.fc3(x)              # output(10)
        return x

需注意：

• pytorch 中 tensor（也就是输入输出层）的 通道排序为：[batch, channel, height, width]
• pytorch中的卷积、池化、输入输出层中参数的含义与位置，可配合下图一起使用：

 1.1 卷积 Conv2d

我们常用的卷积（Conv2d）在pytorch中对应的函数是：

torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')

一般使用时关注以下几个参数即可：

• in_channels：输入特征矩阵的深度。如输入一张RGB彩色图像，那in_channels=3
• out_channels：输入特征矩阵的深度。也等于卷积核的个数，使用n个卷积核输出的特征矩阵深度就是n
• kernel_size：卷积核的尺寸。可以是int类型，如3 代表卷积核的height=width=3，也可以是tuple类型如(3, 5)代表卷积核的height=3，width=5
• stride：卷积核的步长。默认为1，和kernel_size一样输入可以是int型，也可以是tuple类型
• padding：补零操作，默认为0。可以为int型如1即补一圈0，如果输入为tuple型如(2, 1) 代表在上下补2行，左右补1列。

 经卷积后的输出层尺寸计算公式为：

                                           

•  输入图片大小 W×W（一般情况下Width=Height）
• Filter大小 F×F
• 步长 S
• padding的像素数 P

1.2 池化 MaxPool2d

最大池化（MaxPool2d）在 pytorch 中对应的函数是：

MaxPool2d(kernel_size, stride)

1.3 Tensor的展平：view()

注意到，在经过第二个池化层后，数据还是一个三维的Tensor (32, 5, 5)，需要先经过展平后(32*5*5)再传到全连接层：

  x = self.pool2(x)            # output(32, 5, 5)
  x = x.view(-1, 32*5*5)       # output(32*5*5)
  x = F.relu(self.fc1(x))      # output(120)

1.4 全连接 Linear

全连接（ Linear）在 pytorch 中对应的函数是：

Linear(in_features, out_features, bias=True)

2. train.py

2.1 导入数据集

导入包

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import time

数据预处理

对输入的图像数据做预处理，即由shape (H x W x C) in the range [0, 255] → shape (C x H x W) in the range [0.0, 1.0]

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

数据集介绍

利用torchvision.datasets函数可以在线导入pytorch中的数据集，包含一些常见的数据集如MNIST等

导入、加载 训练集

# 导入50000张训练图片
train_set = torchvision.datasets.CIFAR10(root='./data', 	 # 数据集存放目录
										 train=True,		 # 表示是数据集中的训练集
                                        download=True,  	 # 第一次运行时为True，下载数据集，下载完成后改为False
                                        transform=transform) # 预处理过程
# 加载训练集，实际过程需要分批次（batch）训练                                        
train_loader = torch.utils.data.DataLoader(train_set, 	  # 导入的训练集
										   batch_size=50, # 每批训练的样本数
                                          shuffle=False,  # 是否打乱训练集
                                          num_workers=0)  # 使用线程数，在windows下设置为0

导入、加载 测试集

# 导入10000张测试图片
test_set = torchvision.datasets.CIFAR10(root='./data', 
										train=False,	# 表示是数据集中的测试集
                                        download=False,transform=transform)
# 加载测试集
test_loader = torch.utils.data.DataLoader(test_set, 
										  batch_size=10000, # 每批用于验证的样本数
										  shuffle=False, num_workers=0)
# 获取测试集中的图像和标签，用于accuracy计算
test_data_iter = iter(test_loader)
test_image, test_label = test_data_iter.next()

2.2 训练过程

名词	定义
epoch	对训练集的全部数据进行一次完整的训练，称为 一次 epoch
batch	由于硬件算力有限，实际训练时将训练集分成多个批次训练，每批数据的大小为 batch_size
iteration 或 step	对一个batch的数据训练的过程称为 一个 iteration 或 step

以本demo为例，训练集一共有50000个样本，batch_size=50，那么完整的训练一次样本：iteration或step=1000，epoch=1

2.3 使用GPU/CPU训练

使用下面语句可以在有GPU时使用GPU，无GPU时使用CPU进行训练

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

也可以直接指定

device = torch.device("cuda")
# 或者
# device = torch.device("cpu")

对应的，需要用to()函数来将Tensor在CPU和GPU之间相互移动，分配到指定的device中计算

net = LeNet()
net.to(device) # 将网络分配到指定的device中
loss_function = nn.CrossEntropyLoss() 
optimizer = optim.Adam(net.parameters(), lr=0.001) 

for epoch in range(5): 

    running_loss = 0.0
    time_start = time.perf_counter()
    for step, data in enumerate(train_loader, start=0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs.to(device))				  # 将inputs分配到指定的device中
        loss = loss_function(outputs, labels.to(device))  # 将labels分配到指定的device中
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if step % 1000 == 999:    
            with torch.no_grad(): 
                outputs = net(test_image.to(device)) # 将test_image分配到指定的device中
                predict_y = torch.max(outputs, dim=1)[1]
                accuracy = (predict_y == test_label.to(device)).sum().item() / test_label.size(0) # 将test_label分配到指定的device中

                print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                      (epoch + 1, step + 1, running_loss / 1000, accuracy))

                print('%f s' % (time.perf_counter() - time_start))
                running_loss = 0.0

print('Finished Training')

save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)

打印信息如下：

cuda
[1,  1000] train_loss: 1.569  test_accuracy: 0.527
18.727597 s
[2,  1000] train_loss: 1.235  test_accuracy: 0.595
17.367685 s
[3,  1000] train_loss: 1.076  test_accuracy: 0.623
17.654908 s
[4,  1000] train_loss: 0.984  test_accuracy: 0.639
17.861825 s
[5,  1000] train_loss: 0.917  test_accuracy: 0.649
17.733115 s
Finished Training

3. predict.py

# 导入包
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import LeNet

# 数据预处理
transform = transforms.Compose(
    [transforms.Resize((32, 32)), # 首先需resize成跟训练集图像一样的大小
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 导入要测试的图像（自己找的，不在数据集中），放在源文件目录下
im = Image.open('horse.jpg')
im = transform(im)  # [C, H, W]
im = torch.unsqueeze(im, dim=0)  # 对数据增加一个新维度，因为tensor的参数是[batch, channel, height, width] 

# 实例化网络，加载训练好的模型参数
net = LeNet()
net.load_state_dict(torch.load('Lenet.pth'))

# 预测
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].data.numpy()
print(classes[int(predict)])

输出即为预测的标签。

其实预测结果也可以用 softmax 表示，输出10个概率：

with torch.no_grad():
    outputs = net(im)
    predict = torch.softmax(outputs, dim=1)
print(predict)

输出结果中最大概率值对应的索引即为 预测标签的索引

tensor([[2.2782e-06, 2.1008e-07, 1.0098e-04, 9.5135e-05, 9.3220e-04, 2.1398e-04,
         3.2954e-08, 9.9865e-01, 2.8895e-08, 2.8820e-07]])