PyTorch入门项目之卷积神经网络实现Fashion-MNIST图像分类（详细注释！）

PyTorch入门项目，卷积神经网络实现图像分类任务

---

前言

PyTorch目前是学界最流行的深度学习框架之一。对于已经有了深度学习基础的学生而言，急需一个小项目学习练手。然而，网上讲授许多Torch视频教程虽称快速，但很多还是废话较多，部分人可能没有这个耐心看下去，所以我记录了我初次使用torch的一个全流程的训练步骤：导入包、数据集准备、网络搭建、训练、测试、准确率打印、模型保存和loss曲线的内容。

---

提示：本文基于PyTorch1.7.1+cu101、Python3.7，仿AlexNet，基于FashionMNIST数据集；（没有版本说明和代码注释的都是耍流氓）

一、PyTorch安装

关于PyTorch的安装，这里不做过多的叙述，方法很多。
首先，根据自己显卡的CUDA Version来选择对应版本：
cmd命令下输入nvidia-smi 查看：

比如，我的是11.2，可以选择支持11.2及以下版本的Torch版本安装。进入Torch官网：https://pytorch.org/选择：

不知道怎么安装cudatoolkit的话，最简单的就是创建好conda环境后(conda环境配置等问题在此不叙述)，进入环境直接复制：

pip3 install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio===0.9.0

（可能会比较慢，建议用镜像）
当然，不建议选最新的，教程比较少，可以点previous version of PyTorch选择稍低一两个版本的。

二、实践步骤

1.引入相关库

代码如下（示例）：（没有相关module的pip install 下）

import torch
import torchvision  # 数据集用
import torch.nn as nn  # 搭建网络
import torch.utils.data as Data  # 加载数据
import time   # 计时用
import matplotlib.pyplot as plt  # 绘图用

2.超参数设置

代码如下（示例）：

# hyper_parameters setting
DOWNLOAD = True  # 'True'表示需要下载数据集，'False' 表示已经下载好.
BATCH_SIZE = 128   #根据自己电脑实际设置 16 32 64等
EPOCH = 10  # 迭代次数
learning_rate = 0.05  #学习率
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 自动选择训练设备，cpu / gpu（如果支持）

根据需要设置超参数，方便更改程序。

3.数据集准备

以FashionMNIST为例，该数据集可看做是MNIST的升级版，是一些服饰的分类。包含10个类，训练集的图片有60000张，测试集10000张，均为28*28的灰度图像。

## load dataset

train_data = torchvision.datasets.FashionMNIST(
 root='./FashionMnist',  # 数据存放目录
 train=True,  # 用于训练
 transform=torchvision.transforms.ToTensor(),  # 转换为[0,1]的tensor
 download=DOWNLOAD, #是否下载数据集
)

test_data = torchvision.datasets.FashionMNIST(
 root='./FashionMnist', 
 train=False, # 用于测试  
 transform=torchvision.transforms.ToTensor(),
 download=DOWNLOAD)

print(train_data[0]) # 打印查看数据情况 是一个tensor 图像数据和标签

# 数据里的特征和标签按batch分别打包
# shuffle 表示是否随机打乱数据  
# num_workers 表示是否开启新线程加速数据的加载，但设为其他数字时windows系统常报错，所以最好设为0

train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

test_loader = Data.DataLoader(dataset=test_data, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

其他数据集的导入只需要换一下torchvision.datasets.xxx和路径就可以了，其他选择可以保持不变。

4.网络搭建

在这里搭建的卷积神经网络为AlexNet，但又与原版本有些不同，由于图片较小，所有的卷积都变成3*3的。共有5个卷积层，3个全连接层（使用dropout）。

class MyCNN(torch.nn.Module):
    def __init__(self, input_channel=1, output_channel=10):
        super(MyCNN, self).__init__()
        self.conv1 = torch.nn.Sequential(
            nn.Conv2d(in_channels=input_channel,out_channels=96,kernel_size=3,stride=1,padding=1),
            nn.BatchNorm2d(96),  # 批标准化
            nn.ReLU(), # relu激活函数
            nn.MaxPool2d(kernel_size=3,stride=2),  # 最大池化层
        )
        self.conv2 = torch.nn.Sequential(
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.conv3 = torch.nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.conv4 = torch.nn.Sequential(
            nn.Conv2d(in_channels=384, out_channels=384, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.conv5 = torch.nn.Sequential(
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.classifier = nn.Sequential(
            nn.Linear(256 * 2 * 2, 2048),  # 这里是根据输入图像大小算出来的，只有池化操作会缩小图像，28*28的图像池化三次后图像大小为2*2
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(2048, output_channel),
        )

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = x.view(x.size(0),-1)  # 变换维度，相当于拉直操作

        output = self.classifier(x)
        
		return output

nn.Sequential() :用于快速搭建网络序列；
input_channel :表示输入该卷积层的通道数目；
output-channel :表示输入通道数，也就是该卷积层有多少个卷积核。

5.训练设置

训练的优化器、loss选择。

cnn = MyCNN(input_channel=1, output_channel=10)  # 创建网络，不能忘
print(cnn) # 打印网络结构
cnn.to(device) # 将网络加载到设备

# optimizer  优化器的选择和损失函数的定义 
# 在此处选择SGD随机梯度下降优化，交叉熵损失
optimizer = torch.optim.SGD(cnn.parameters(), lr=learning_rate)
loss_func = nn.CrossEntropyLoss()

6.训练过程

代码如下（示例）：

# 为绘图、保存和输出做些数据准备
train_epochs_loss = []
train_acc = []
test_epochs_loss = []
test_acc = []
best_acc = 0

# training and testing
print('Training and Testing ...\n')
for epoch in range(EPOCH):
    cnn.train()  # 进入训练模式，养成习惯
    train_l_sum, train_acc_sum, n, start = 0.0, 0.0, 0, time.time()
    for step, (tr_x, tr_y) in enumerate(train_loader):
        tr_x, tr_y = tr_x.to(device), tr_y.to(device)
        output = cnn(tr_x)  # 网络训练好的输出
        loss = loss_func(output, tr_y)  # 计算损失
        optimizer.zero_grad() # 梯度先置0
        loss.backward()  # 根据损失函数后向传播，计算各参数的梯度
        optimizer.step()  # 参数更新

        train_l_sum += loss.item()  # 累计总损失
        train_acc_sum += (output.argmax(dim=1) == tr_y).sum().item()  # 累计预测正确的数目
        n += tr_y.shape[0]

    train_epochs_loss.append((train_l_sum/(step+1)))
    train_acc.append(train_acc_sum / n)

7.测试过程

代码如下（示例）：跟训练时类似，只是不需要后向传播更新参数，

    cnn.eval()  # 测试形态，不能忘
    acc_sum, num, te_loss = 0.0, 0, 0.0
    with torch.no_grad():  # 这一句也很重要，有时可以减少内存占用（具体原因未知）
        for bat_idx, (X, y) in enumerate(test_loader):
            te_out = cnn(X.to(device))
            acc_sum += (te_out.argmax(dim=1) == y.to(device)).sum().item()
            te_loss += loss_func(te_out, y.to(device)).item()
            num += y.shape[0]
    test_acc_epoch = acc_sum / num
    test_loss = te_loss / (bat_idx+1)
    test_acc.append(test_acc_epoch)
    # 打印训练和测试，查看本次epoch结果
    print('epoch: %d, loss %.4f, train_acc: %.3f,  test_acc: %.3f,  test_loss: %.3f, time %.1f sec'
        % (epoch + 1, train_l_sum / (step+1), train_acc_sum / n,  test_acc_epoch, test_loss, time.time() - start))
    test_epochs_loss.append(test_loss)

8.保存模型

代码如下（示例）：根据测试准确率来保存最优模型

       if test_acc_epoch > best_acc:
        torch.save(cnn.state_dict(),'cnn_params_best.pkl')
        best_acc = test_acc_epoch
        print('cnn has saved\n')

这里采用了torch.save(cnn.state_dict(),'cnn_params_best.pkl')，只保存网络参数，当加载时需要先搭建好网络结构，

cnn = MyCNN(input_channel=1, output_channel=10)
cnn.to(device)

然后再加载：
cnn.load_state_dict(torch.load('cnn_params_best.pkl'))

另外一种保存网络的方法是保存整个网络，包括结构
torch.save(cnn,'net.pkl')
加载时直接赋值
cnn = torch.load('net.pkl')

9.画出accuracy and loss曲线

plt.subplot(121)
plt.plot(train_acc[:],'-o',label="train_acc")
plt.plot(test_acc[:],'-o',label="test_acc")
plt.title('epochs_accuracy')
plt.legend()
plt.subplot(122)
plt.plot(train_epochs_loss[:],'-o',label="train_loss")
plt.plot(test_epochs_loss[:],'-o',label="test_loss")
plt.title("epochs_loss")
plt.legend()
plt.savefig('eqochs_acc_loss.png')
plt.show()

10.结果

训练过程中的输出：

最后的训练曲线：（赶时间，我只让它迭代了5次，，）

三、总结

PyTorch还是比较好入门的，只不过这也只是一个基础的入门，一些复杂的网络结构和输出还要等待实践！
（将各部分复制一下即可完美运行，我重新复制到编辑器里亲测了~）

另外，TensorFlow版的经典卷积网络代码在这里，感谢支持！代码若有问题欢迎指正交流。