AlexNet代码讲解（Pytorch）

引言

文通过代码实现了AlexNet算法，使用的是pytorch框架，版本为1.7.1。另外本专栏的所有算法都有对应的Libtorch版本（Libtorch版本的AlexNet地址），算法原理本文不做过多阐述。本文针对小白对代码以及相关函数进行讲解，建议配合代码进行阅读，代码中我进行了详细的注释，因此读者可以更加容易理解代码的含义，本文只展示了部分代码，全部代码可以通过GitHub下载。

本文使用的数据集为0~9的手写数据集，全部代码主要分为以下几个部分：

1、定义AlexNet网络结构（model.py）

2、训练卷积神经网络（train.py）

3、输入图片预测结果（predict.py）

4、将保存的模型参数 .pth文件(Pythorch用)转变为 .pt文件(Libtorch用) （pytocpp.py）

model.py

AlexNet的网络结构如下所示：

首先是特征提取部分的网络结构，其中每一次卷积后都需要加ReLu激活函数。

此部分代码如下

 self.features = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(96, 256, 5, 1, 2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(256, 384, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3, 2)
        )

然后是线性分类部分的网络结构：

此部分代码如下：

self.classifier = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(256*6*6, 2048),
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, NUM_CLASS) # NUM_CLASS为类别总数
        )

先介绍以下用到的函数吧

import torch.nn as nn
import torch
# 卷积层  计算公式为 -> 卷积后图像尺寸=(原图像尺寸+2*填充大小-卷积核大小)/步长+1
# in_channels为输入通道数
# out_channels为输出通道数
# kernel_size为卷积核尺寸 比如11代表（11*11）
# stride卷积的步长
# padding为零填充数
nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4, padding=2)
 
# 最大池化层
# kernel_size为卷积核大小
# stride为步长
nn.MaxPool2d(kernel_size=3, stride=2)
 
# ReLu激活函数
# inplace=True会改变输入数据的值,节省反复申请与释放内存的空间与时间,效率更好
nn.ReLU(inplace=True)
 
nn.Dropout(0.5),  # Dropout随机丢弃神经元，0.5代表随机丢弃50%
 
torch.flatten(x, start_dim=1) # （将矩阵x展开成一维行向量）
 
# 全连接层
# in_features 输入部分神经元个数
# out_features 输出部分神经元个数
nn.Linear(in_features,out_features)

以0~9的手写数据集为例，在AlexNet中要求输入224*224*3的彩色三通道RGB图片,经过特征提取部分的神经网络输出256*6*6大小的张量，再经过线性分类部分的神经网络得到1*10的张量，假设这10个数中第4个数最大，则最终预测此图像为数字3。至于权重初始化代码，是官方例程中给出的，参考一下就好。

train.py

定义完网络结构之后就可以正式开始训练了！

具体分为以下几个步骤：

1、定义数据集

2、定义网络结构

3、定义损失函数以及优化器

4、开始训练

首先是数据集的定义，本文用到了torchvision中datasets.ImageFolder函数，此函数对数据集的摆放格式有一定要求，以0~9的手写数据集为例，具体布置格式如下图所示：

在定义数据集的同时需要对所有的图片进行预处理，包括将图片resize成（224，224）的大小，转变为张量，进行标准化等等。

部分代码如下所示：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
 
# transforms.Compose可以对张量进行一系列操作，各种操作存储在列表内
transform = {
# ToTensor()能够把像素的值域范围从0-255变换到0-1之间，
# 而后面的transform.Normalize()则把0-1变换到(-1,1).
    'train': transforms.Compose([transforms.Resize((224, 224)),
                                 transforms.RandomHorizontalFlip(),  # 随机翻转
                                 transforms.ToTensor(),  # 转换为张量
                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化处理
                                 ]),
    'test': transforms.Compose([transforms.Resize((224, 224)),
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
}
train_dataset = datasets.ImageFolder(root='./AlexNet/Project1/DATASET/TRAIN', transform=transform['train'])
test_dataset = datasets.ImageFolder(root='./AlexNet/Project1/DATASET/TEST', transform=transform['test'])
 
# 参数shuffle代表是否打乱数据集（以乱序排列）
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False)

定义完数据集后，接下来是定义网络结构：（如果使用GPU进行训练，需要设置为CUDA），代码如下：

# 2、定义网络结构并设置为CUDA
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 这个AlexNet是我们在model.py中定义的类（网络结构）
net = AlexNet(NUM_CLASS=10, init_weight=True)
net.to(device)  # 转化为CUDA

接下来是定义损失函数以及优化器，我们这里使用的是交叉熵损失函数(CrossEntropyLoss)以及SGD优化器(随机梯度下降)，如果使用GPU进行训练损失函数也要设置为CUDA，具体代码如下：

# 3、定义损失函数及优化器，损失函数设置为CUDA
loss_function = nn.CrossEntropyLoss()
optimizer = optioms.SGD(params=net.parameters(), lr=learning_rate)
loss_function.to(device)

好的，最后就开始训练了，在训练过程中学习率(learning rate)逐步减小会比较好，更改学习率则需要对优化器(optimizer)进行设置，在optimizer中存有字典(param_groups)，在字典中可以设置学习率,设置方法如下：

for param_group in optimizer.param_groups:   # 其中的元素是2个字典；optimizer.param_groups[0]： 长度为6的字典，包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数；
                                                # optimizer.param_groups[1]： 好像是表示优化器的状态的一个字典；
    param_group['lr'] = learning_rate      # 更改全部的学习率

然后就需要从上文定义的train_loader中取数据(image与target)进行训练（若使用GPU训练则image与target都要设置为CUDA），具体步骤为：1、将图片输入网络得到1*10的张。2、将神经网络的输出与标签（target）进行对比并计算其损失。3、通过优化器进行梯度下降和反向传播更新参数。4、保存模型参数。5、不断重复（迭代）以上步骤直到达到要求（损失值小于设定值或者完成设定的迭代次数）。部分代码如下所示。

# 4、开始训练
for epoch in range(num_epochs):
    net.train()  # 网络有Dropout，BatchNorm层时一定要加
    if epoch == 4:
        learning_rate = 0.0001  # 第四次迭代时学习率设置为0.0001
    if epoch == 6:
        learning_rate = 0.00001
    for param_group in optimizer.param_groups:   # 其中的元素是2个字典；optimizer.param_groups[0]： 长度为6的字典，包括[‘amsgrad’, ‘params’, ‘lr’, ‘betas’, ‘weight_decay’, ‘eps’]这6个参数；
                                                # optimizer.param_groups[1]： 好像是表示优化器的状态的一个字典；
        param_group['lr'] = learning_rate      # 更改全部的学习率
    print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
    print('Learning Rate for this epoch: {}'.format(learning_rate))
 
    total_loss = 0.
    for i, (images, target) in enumerate(train_loader):  # image为图片，target为其对应的标签（类别名）
        images, target = images.cuda(), target.cuda()  # 设置为CUDA
        pred = net(images)    # 图片输入网络得到预测结果
        loss = loss_function(pred, target)  # 将预测结果与实际标签比对（计算两者之间的损失值）
        total_loss += loss.item()
 
        optimizer.zero_grad()   # 将梯度归零，有助于梯度下降
        loss.backward()    # 反向传播 计算梯度
        optimizer.step()   # 根据梯度 更新模型参数
        if (i + 1) % 5 == 0: # 打印训练的信息
            print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f' % (epoch + 1, num_epochs,
                                                                                 i + 1, len(train_loader), loss.item(), total_loss / (i + 1)))
    validation_loss = 0.0
    net.eval()
    for i, (images, target) in enumerate(test_loader):  # 导入dataloader 说明开始训练了  enumerate 建立一个迭代序列
        images, target = images.cuda(), target.cuda()
        pred = net(images)    # 将图片输入
        loss = loss_function(pred, target)
        validation_loss += loss.item()   # 累加loss值  （固定搭配）
    validation_loss /= len(test_loader)  # 计算平均loss
    if best_test_loss > validation_loss:
        best_test_loss = validation_loss
        print('get best test loss %.5f' % best_test_loss)
        torch.save(net.state_dict(), 'AlexNet.pth')  # 保存模型参数

predict.py

知道了如何训练那预测阶段也十分容易了，大致思路为，使用opencv读取一张图片然后经过和训练阶段一样的预处理操作，将其输入神经网络得到1*10维的张量，通过判断10个数中第几个数最大，就能知道该图片所属的类别。废话不多说，代码如下

import cv2
import torch
 
transform = transforms.Compose([ transforms.ToTensor(),  # 转换为张量
                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
 
# 进行预测
class Predict:
    def __init__(self, image_root, net1): # image_root->需要预测的图片路径，net1->网络结构
        self.img_root = image_root
        self.model = net1
 
    def result(self):
        img = cv2.imread(img_root)  # opencv读取图片
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # BGR->RGB
        img = cv2.resize(img, (224, 224))  # 将图片尺寸变为224*224
        img = transform(img)   # 将图片对应的像素矩阵变为张量并且标准化  size=(3,224,224))
        img = torch.unsqueeze(img, 0)  # 执行完后尺寸为(1,3,224,224)
        result = self.model(img)  # 将图片输入神经网络得到结果
        return result

注意事项：

通过opencv读取到的图像彩色三通道的BGR格式的图像，而上文我们说到AlexNet需要输入的是RGB格式的图片，所以需要执行img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)这段代码，值得注意的是 transforms.ToTensor() 这段代码 将尺寸为 （224，224，3）的图像矩阵变成了 （3，224，224）。在Pytorch中神经网络处理图片的格式应为 [ batch_size , channel , height , width ] ,所以需要通过torch.unsqueeze来实现这个要求，另外这个操作在训练时是由 train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True) 这行代码实现的，batch_size指的是同时输入神经网络的图片的个数，batch_size为1则输出(1,10)的张量得到一张图片的结果，batch_size为2则会输出(2,10)的张量得到两张图片识别的结果。

另外在此文件中我还附上了绘制混淆矩阵的方法，只需要将代码83行之后的注释全部清除即可（运行会花费一定时间）。

最后附上效果图：

结尾

由于本专栏所有代码都有C++版本，所以提供了一段代码用于将Python版本的模型参数文件转变为C++版本的，如果需要请注意Pytorch与Libtorch版本要一致（在这里是都为1.7.1），读者也可以使用C++进行训练。本文的代码也可以使用其他数据集进行训练，但是需要对代码进行小小的更改，更改方法放在了GitHub中，另外本文用到的手写数据集也可以一并在GitHub中下载。