AlexNet示例代码与解析
- 1、前言
- 2、模型tips
- 3、模型架构
- 4、模型代码
- backbone
- train
- predict
- 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=17,
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模型可视化架构图: