【深度学习|基础算法】2.AlexNet（附训练|推理代码）

本文链接：https://blog.csdn.net/ycx_ccc/article/details/137067799

AlexNet示例代码与解析

1、前言
2、模型tips
3、模型架构
4、模型代码
5、模型训练
6、导出onnx模型

1、前言

AlexNet由Hinton和他的学生Alex Krizhevsky设计，模型名字来源于论文第一作者的姓名Alex。该模型以很大的优势获得了2012年ISLVRC竞赛的冠军网络，分类准确率由传统的 70%+提升到 80%+，自那年之后，深度学习开始迅速发展。
ImageNet是一个在2009年创建的图像数据集，从2010年开始到2017年举办了七届的ImageNet 挑战赛——ImageNet Large Scale Visual Recognition ChallengeI (LSVRC)，在这个挑战赛上诞生了AlexNet、ZFNet、OverFeat、VGG、Inception、ResNet、WideResNet、FractalNet、DenseNet、ResNeXt、DPN、SENet 等经典模型。
摘录：CNN经典网络模型（二）：AlexNet简介及代码实现（PyTorch超详细注释版）

2、模型tips

使用了多GPU训练，当时的硬件资源有限，一块显卡的显存并不能满足AlexNet的训练，因此作者将其分成了两部分，每一块显卡负责一部分的特征图，最后在全连接层的时候，将每块显卡提取的特征拼接成一个，当时使用了分组卷积来完成这个需求，后来证实了分组卷积能够达到普通卷积的精度。
使用了ReLu作为激活函数来训练模型。
使用了局部响应归一化来对网络层进行归一化，目的是抑制反馈较小的神经元的值，增大反馈明显的神经元的值。（在VGG的论文中指出，LRN并灭有什么明显的效果，新的深度学习网络中已使用其他归一化方法来代替LRN)。
使用了覆盖池化，覆盖池化的意思即当池化的步长小于池化核时，相邻的池化核之间会互相重叠。论文中指出这种池化方式可以缓解过拟合。
使用了dropout来缓解过拟合。dropout即在训练的过程中随机的让一些隐层中的节点置为0（在本轮训练中不参与前向传播与反向传播），因此，在每一轮训练的时候，模型都会随机的得到一个与上一论训练中不太一样的网络架构，这样的做法可以降低神经元之间的依赖性（耦合性），使每个神经元能够在一次次的训练中学习到更为可靠的特征。尽管dropout会降低训练收敛的速度，但是可以有效的缓解过拟合程度。

3、模型架构

在这里插入图片描述

4、模型代码

我们将使用开源的深度学习框架PyTorch来搭建并训练我们的模型

backbone

import torch
import torchvision
from torch import nn

class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        #特征提取层
        self.CONV = nn.Sequential(
            nn.Conv2d(3, 96, 11, 4, 2),     # 1 * 224 * 224 * 3 ---> 1 * 55 * 55 * 96
            nn.ReLU(),
            nn.MaxPool2d(3, 2),             # 1 * 55 * 55 * 96 ---> 1 * 27 * 27 * 96

            nn.Conv2d(96, 256, 5, 1, 2),    # 1 * 27 * 27 * 96 ---> 1 * 27 * 27 * 256
            nn.ReLU(),
            nn.MaxPool2d(3, 2),             # 1 * 27 * 27 * 256 ---> 1 * 13 * 13 * 256

            nn.Conv2d(256, 384, 3, 1, 1),   # 1 * 13 * 13 * 256 ---> 1 * 13 * 13 * 384
            nn.ReLU(),

            nn.Conv2d(384, 384, 3, 1, 1),   # 1 * 13 * 13 * 384---> 1 * 13 * 13 * 384
            nn.ReLU(),

            nn.Conv2d(384, 256, 3, 1, 1),   # 1 * 13 * 13 * 384 ---> 1 * 13 * 13 * 256
            nn.ReLU(),
            nn.MaxPool2d(3, 2),             # 1 * 13 * 13 * 256 ---> 1 * 6 * 6 * 256
        )

        #将多维的张量进行平坦化处理
        #默认从第一个维度到最后一个维度拼接
        self.flatten = nn.Flatten()         # 1 * 6 * 6 * 256 ---> 1 * 9216

        #全连接层
        self.FC = nn.Sequential(
            # 全连接层1
            nn.Linear(in_features=6 * 6 * 256, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            # 全连接层2
            nn.Linear(in_features=4096, out_features=4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            # 全连接层3
            nn.Linear(in_features=4096, out_features=1000),
            nn.Dropout(0.5),
            # 全连接层4 分类层，几分类就是降维到几
            nn.Linear(in_features=1000, out_features=10),
        )

    def forward(self, x):
        x = self.CONV(x)
        x = self.flatten(x)
        x = self.FC(x)
        return x

if __name__ == "__main__":
    x = torch.randn([1,3,224,224])
    model = AlexNet()
    y = model(x)
    print(x)

注意这里的输入是224 * 224的宽高，有的输入是227 * 227的宽高，可以根据不同的输入来计算对应的卷积与池化的参数，对最终的结果几乎没有影响。

搭建模型时用到的torch库的api：
nn.Sequential()： 可以允许将整个容器视为单个模块（即相当于把多个模块封装成一个模块），forward()方法接收输入之后，nn.Sequential()按照内部模块的顺序自动依次计算并输出结果。
nn.Conv2d()：in_channel，out_channel，kernel_size，stride，padding。
nn.ReLu()：ReLu激活函数。
nn.MaxPool2dkernel_size，stride。
nn.Flatten()：将张量扁平化处理，默认从第一个维度到最后一个维度的所有信息进行一维处理。start_dim=1，end_dim=-1。
nn.Linear()：线性层，或者叫全连接层。in_features，out_features，is_bias(bool)。
nn.Dropout()：dropout方法，随机失活神经元节点。p（失活节点的占比），in_place（是否改变输入数据），dropout只在训练阶段开启，在推理阶段不开启。

train

import torch
import torch.nn as nn
from AlexNet import AlexNet
from torch.optim import lr_scheduler
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import os
import matplotlib.pyplot as plt

# 解决中文显示问题
# 运行配置参数中的字体（font）为黑体（SimHei）
plt.rcParams['font.sans-serif'] = ['simHei']
# 运行配置参数总的轴（axes）正常显示正负号（minus）
plt.rcParams['axes.unicode_minus'] = False

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

ROOT_TRAIN = 'dataset'
ROOT_TEST = 'dataset'

normalize = transforms.Normalize(
    [0.5, 0.5, 0.5],
    [0.5, 0.5, 0.5]
)

train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    #以0.5的概率来竖直翻转给定的PIL图像
    transforms.RandomVerticalFlip(),
    transforms.ToTensor(),
    normalize,
])

val_transform = transforms.Compose([
    transforms.Resize((224,224)),
    transforms.ToTensor(),
    normalize,
])

#加载训练数据集
train_datasets = datasets.ImageFolder(ROOT_TRAIN, transform=train_transform)
train_dataloader = DataLoader(train_datasets, batch_size=32, shuffle=True)

val_datasets = datasets.ImageFolder(ROOT_TEST, transform=val_transform)
val_dataloader = DataLoader(val_datasets, batch_size=32, shuffle=True)

#实例化模型对象
model = AlexNet().to(device)

#定义交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()

#定义优化器
optimizer_ = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

#学习率每十轮降低为之前的0.5
lr_scheduler = lr_scheduler.StepLR(optimizer_, step_size=10, gamma=0.5)

#定义训练函数
def train(dataloader, model, loss_fn, optimizer):
    loss, current, n = 0.0, 0.0, 0.0
    #batch为索引， x和y分别是图片和类别
    for batch, (x, y) in enumerate(dataloader):
        #前向传播
        image, y = x.to(device), y.to(device)
        output = model(image)
        curr_loss = loss_fn(output, y)
        _, pred = torch.max(output, dim=1)
        #计算每个批次的准确率
        curr_acc = torch.sum(y == pred)/output.shape[0]

        #反向传播
        #清空之前的梯度
        optimizer.zero_grad()
        #计算当前的梯度
        curr_loss.backward()
        #根据梯度更新网络参数
        optimizer.step()

        #损失叠加
        loss += curr_loss.item()
        #精度叠加
        current += curr_acc.item()
        n = n + 1

    #训练的平均损失和平均精度
    train_loss = loss / n
    train_acc = current / n
    print('train loss = ' + str(train_loss))
    print('train accuracy = ' + str(train_acc))
    return train_loss, train_acc

#定义验证函数
def val(dataloader, model, loss_fn):
    loss, current, n = 0.0, 0.0, 0.0
    #eval()：如果模型中存在BN和dropout则不启用，以防改变权值
    model.eval()
    with torch.no_grad():
        for batch, (x, y) in enumerate(dataloader):
            #前向传播
            image, y = x.to(device), y.to(device)
            output = model(image)
            curr_loss = loss_fn(output, y)
            _, pred = torch.max(output, dim=1)
            curr_acc = torch.sum(y == pred) / output.shape[0]
            loss += curr_loss.item()
            current += curr_acc.item()
            n = n + 1

    val_loss = loss / n
    val_acc = current / n
    print('val loss = ' + str(val_loss))
    print('val accuracy = ' + str(val_acc))
    return val_loss, val_acc

#定义画图函数
def plot_loss(train_loss, val_loss):
    plt.plot(train_loss, label='train loss')
    plt.plot(val_loss, label='val loss')

    plt.legend(loc='best')
    plt.xlabel('loss')
    plt.ylabel('epoch')
    plt.title("训练集和验证集的loss值对比图")
    plt.show()

def plot_acc(train_acc, val_acc):
    plt.plot(train_acc, label='train acc')
    plt.plot(val_acc, label='val acc')

    plt.legend(loc='best')
    plt.xlabel('acc')
    plt.ylabel('epoch')
    plt.title("训练集和验证集的acc值对比图")
    plt.show()

#开始训练
loss_train = []
acc_train = []
loss_val = []
acc_val = []

#训练次数
epoch = 200
#用于判断什么时候保存模型
min_acc = 0
for t in range(epoch):
    # lr_scheduler.step()
    print(f"epoch{t+1}-------------------------------")
    #训练模型
    train_loss, train_acc = train(train_dataloader, model, loss_fn, optimizer_)
    #验证模型
    val_loss, val_acc = val(val_dataloader, model, loss_fn)
    print("\n")
    loss_train.append(train_loss)
    acc_train.append(train_acc)
    loss_val.append(val_loss)
    acc_val.append(val_acc)
    folder = 'save_model'

    # 保存最好的模型权重
    if val_acc > min_acc:
        if not os.path.exists(folder):
            os.mkdir(folder)
        min_acc = val_acc
        torch.save(model.state_dict(), f"{folder}/model_best.pth")

    if t == epoch - 1:
        torch.save(model.state_dict(), f"{folder}/model_last.pth")
        print("=============训练完毕==============\n" + f"best pth saved as {folder}/model_best.pth\n" + f"last pth saved as {folder}/model_last.pth\n")

plot_loss(loss_train, loss_val)
plot_acc(acc_train, acc_val)

predict

import torch
from AlexNet import AlexNet
from torch.autograd import Variable
from torchvision import transforms
from torchvision.transforms import ToPILImage
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

# ROOT_TRAIN = 'D:/pycharm/AlexNet/data/train'
ROOT_TEST = 'dataset'

# 将图像的像素值归一化到[-1,1]之间
normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    normalize
])

# 加载训练数据集
val_dataset = ImageFolder(ROOT_TEST, transform=val_transform)

# 如果有NVIDA显卡，转到GPU训练，否则用CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 模型实例化，将模型转到device
model = AlexNet().to(device)

# 加载train.py里训练好的模型
model.load_state_dict(torch.load(r'save_model/model_best.pth'))

# 结果类型
classes = [
    "cat",
    "dog"
]

# 把Tensor转化为图片，方便可视化
show = ToPILImage()

# 进入验证阶段
model.eval()
for i in range(10):
    x, y = val_dataset[i][0], val_dataset[i][1]
    # show()：显示图片
    show(x).show()
    # torch.unsqueeze(input, dim)，input(Tensor)：输入张量，dim (int)：插入维度的索引，最终扩展张量维度为4维
    x = Variable(torch.unsqueeze(x, dim=0).float(), requires_grad=False).to(device)
    with torch.no_grad():
        pred = model(x)
        # argmax(input)：返回指定维度最大值的序号
        # 得到预测类别中最高的那一类，再把最高的这一类对应classes中的那一类
        predicted, actual = classes[torch.argmax(pred[0])], classes[y]
        # 输出预测值与真实值
        print(f'predicted:"{predicted}", actual:"{actual}"')

5、模型训练

猫狗训练集训练二分类模型，生成pth模型。

在这里插入图片描述

6、导出onnx模型

from AlexNet import AlexNet
import torch
import onnx
import os

def export_onnx(pt_path:str, onnx_path:str):
    model = AlexNet()

    model.load_state_dict(torch.load(pt_path))

    dummy_inputs = {
        "input": torch.randn(1, 3, 224, 224, dtype=torch.float),
    }

    output_names = {
        "classes"
    }

    if not os.path.exists(os.path.dirname(onnx_path)):
        os.makedirs(os.path.dirname(onnx_path))

    with open(onnx_path, "wb") as f:
        print(f"Exporting onnx model to {onnx_path}...")
        torch.onnx.export(
            model,
            tuple(dummy_inputs.values()),
            f,
            export_params=True,
            verbose=False,
            opset_version=11,
            do_constant_folding=True,
            input_names=list(dummy_inputs.keys()),
            output_names=output_names,
            # dynamic_axes=dynamic_axes,
        )

if __name__ == "__main__":
    pt_path = "save_model/model_best.pth"
    onnx_path = "save_model/model_best.onnx"
    export_onnx(pt_path, onnx_path)

onnx模型可视化架构图：

在这里插入图片描述

onnxruntime推理

这里有几个需要注意的点：

这里我读取图像用的是opencv，因此读取的图像的通道是BGR排列的，而在predict的前处理中用的是torchvision.datasets.ImageFolder来加载图片的，默认的加载方式是PIL，因此需要用cvtcolor将BGR转成RGB，当然也可以像我一样不使用cvtcolor
这里我用的transform来进行前处理，compose中包含了resize，totensor还有normalize，如果不将ndarray格式的input转成PIL格式的话，transforms.compose中需要先做ToTensor()，然后再进行resize和normalize，否则会报数据类型错误，因为transforms.Resize()无法直接处理ndarray的图，调换顺序之后输入会有些许的差别（因为ToTensor和Resize都有自己的处理方式，顺序不同会产生不同的结果）
onnxruntime的输入是ndarray类型的，因此处理完图像数据之后是tensor的格式，我们需要将其变成ndarray再输入模型中推理。

import numpy as np
import onnxruntime
import torch
from PIL import Image
from torchvision import transforms
import cv2
from torch.autograd import Variable

classes = [
    "cat",
    "dog",
]

mean = [0.5,0.5,0.5]
std = [0.5,0.5,0.5]

def onnx_infer(onnx_path:str, input):
    input = input[:,:,(2,1,0)]
    # ndarray -> PIL
    image = Image.fromarray(input)

    normalize = transforms.Normalize([0.5,0.5,0.5],[0.5,0.5,0.5])
    transformers = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        normalize,
    ])

    # input preprocess
    # ndarray -> tensor
    input = transformers(image)

    input_blob = Variable(torch.unsqueeze(input, dim=0).float(), requires_grad=False)
    input_blob = input_blob.detach().numpy()
    # session 构造
    model_session = onnxruntime.InferenceSession(onnx_path, provider_options=["CUDAExecutionProvider"])

    # 执行推理
    pred = model_session.run(None, {'input':input_blob})
    pred_ = pred[0][0]

    class_ = classes[np.argmax(pred_)]

    print(class_)

if __name__ == "__main__":
    onnx_path = "save_model/model_best.onnx"
    image_path = "cat.0.jpg"
    input_image = cv2.imread(image_path)
    # input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
    onnx_infer(onnx_path, input_image)