本来，前辈们用caffe搭建了一个squeezenet的工程，用起来也还行，但考虑到caffe的停更后续转trt应用在工程上时可能会有版本的问题所以搭建了一个pytorch版本的。
以下的环境搭建不再细说，主要就是pyorch，其余的需要什么pip install什么。
网络搭建
squeezenet的网络结构及其具体的参数如下：
后续对着这张表进行查看每层的输出时偶然发现这张表有问题，一张224×224的图片经过7×7步长为2的卷积层时输出应该是109×109才对，而不是这个111×111。所以此处我猜测要不是卷积核的参数有问题，要不就是这个输出结果有问题。我对了下下面的结果，发现都是从这个111×111的结果得出来的，这个结果没问题；但是我又对了下原有caffe版本的第一个卷积层用的就是这个7×7/2的参数，卷积核也没问题。这就有点矛盾了…这张表出自作者原论文，论文也是发表在顶会上，按道理应该不会有错才对。才疏学浅，希望大家有知道咋回事的能告诉我一声，这里我就还是用这个卷积核的参数了。
squeezenet有以上三个版本，我对了下发现前辈用的是中间这个带有简单残差的结构，为了进行对比这里也就用这个结构进行搭建了。
如下为网络结构的代码：
import torch
import torch.nn as nn
class Fire(nn.Module):
def __init__(self, in_channel, squzee_channel, out_channel):
super().__init__()
self.squeeze = nn.Sequential(
nn.Conv2d(in_channel, squzee_channel, 1),
nn.ReLU(inplace=True)
)
self.expand_1x1 = nn.Sequential(
nn.Conv2d(squzee_channel, out_channel, 1),
nn.ReLU(inplace=True)
)
self.expand_3x3 = nn.Sequential(
nn.Conv2d(squzee_channel, out_channel, 3, padding=1),
nn.ReLU(inplace=True)
)
def forward(self, x):
x = self.squeeze(x)
x = torch.cat([
self.expand_1x1(x),
self.expand_3x3(x)
], 1)
return x
class SqueezeNet_caffe(nn.Module):
"""mobile net with simple bypass"""
def __init__(self, class_num=5):
super().__init__()
self.stem = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=96, kernel_size=7, stride=2),
nn.ReLU(inplace=True),
nn.MaxPool2d(3, 2, ceil_mode=True)
)
self.fire2 = Fire(96, 16, 64)
self.fire3 = Fire(128, 16, 64)
self.fire4 = Fire(128, 32, 128)
self.fire5 = Fire(256, 32, 128)
self.fire6 = Fire(256, 48, 192)
self.fire7 = Fire(384, 48, 192)
self.fire8 = Fire(384, 64, 256)
self.fire9 = Fire(512, 64, 256)
self.maxpool = nn.MaxPool2d(3, 2, ceil_mode=True)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Conv2d(512, class_num, kernel_size=1),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d((1, 1))
)
def forward(self, x):
x = self.stem(x)
f2 = self.fire2(x)
f3 = self.fire3(f2) + f2
f4 = self.fire4(f3)
f4 = self.maxpool(f4)
f5 = self.fire5(f4) + f4
f6 = self.fire6(f5)
f7 = self.fire7(f6) + f6
f8 = self.fire8(f7)
f8 = self.maxpool(f8)
f9 = self.fire9(f8) + f8
x = self.classifier(f9)
x = x.view(x.size(0), -1)
return x
def squeezenet_caffe(class_num=5):
return SqueezeNet_caffe(class_num=class_num)
然后其余的整个工程代码就是pytorch搭建dataset、dataloader，每轮的前向、计算loss、反向传播等都是一个差不多的套路，就不在这里码出来了，直接放上链接，大家有需要可以直接下载（里面也集成了其他的分类网络）。
数据处理
dataset我用的是torchvision.datasets.ImageFolder，所以用目录名称作为数据集的label，目录结构如下：
将每一类的图片都放在对应的目录中，验证集以及测试集的数据集也是按照这样的格式。
运行命令
训练命令：
python train.py -net squeezenet_caffe -gpu -b 64 -t_data 训练集路径 -v_data 验证集路径 -imgsz 100
-net后面跟着是网络类型，都集成了如下的分类网络：
如果有n卡则-gpu使用gpu训练，-b是batch size，-imgsz是数据的input尺寸即resize的尺寸。
测试命令：
python test.py -net squeezenet_caffe -weights 训练好的模型路径 -gpu -b 64 -data 测试集路径 -imgsz 100
出现问题
一开始进行训练一切正常，到后面却忽然画风突变：
loss忽然大幅度上升，acc也同一时刻大幅度下降，然后数值不变呈斜率为0的一条直线。估计是梯度爆炸了（也是到这一步我先从网络结构找原因，对本文的第一张表一层一层对参数和结果才发现表中的问题），网络结构对完没问题，于是打印每个batch的梯度，顺便使用clip进行剪枝限定其最大阈值。
optimizer.zero_grad()
outputs = net(images)
loss = loss_function(outputs, labels)
loss.backward()
grad_max = 0
grad_min = 10
for p in net.parameters():
# 打印每个梯度的模,发现打印太多了一直刷屏所以改为下面的print最大最小值
# print(p.grad.norm())
gvalue = p.grad.norm()
if gvalue > grad_max:
grad_max = gvalue
if gvalue < grad_min:
grad_min = gvalue
print("grad_max:")
print(grad_max)
print("grad_min:")
print(grad_min)
# 将梯度的模clip到小于10的范围
torch.nn.utils.clip_grad_norm(p,10)
optimizer.step()
按道理来说应该会有所改善，但结果是，训练几轮之后依旧出现这个问题。
但是，果然梯度在曲线异常的时候数值也是异常的：
刚开始正常学习的时候梯度值基本上都在e-1数量级的，曲线异常阶段梯度值都如图所示无限接近0,难怪不学习。
我们此时看一下tensorboard，我将梯度的最大最小值write进去，方便追踪：
可以发现在突变处梯度值忽然爆炸激增，猜测原因很可能是学习率太大了，动量振动幅度太大了跳出去跳不回来了。查看设置的学习率超参发现初始值果然太大了（0.1），于是改为0.01。再次运行后发现查看其tensorboard：
这回是正常的了。
但其实我放大查看了梯度爆炸点的梯度值：
发现其最大值没超过10，所以我上面的clip没起到作用，我如果将阈值改成2，结果如下：
发现起到了作用，但曲线没那么平滑，可能改成1或者再小一些效果会更好。但我觉得还是直接改学习率一劳永逸比较简单。
Pytorch模型转TensorRT模型
在训练了神经网络之后，TensorRT可以对网络进行压缩、优化以及运行时部署，并且没有框架的开销。TensorRT通过combines
layers，kernel优化选择，以及根据指定的精度执行归一化和转换成最优的matrix math方法，改善网络的延迟、吞吐量以及效率。
总之，通俗来说，就是训练的模型转trt后可以在n卡上高效推理，对于实际工程应用更加有优势。
首先将pth转onnx：
# pth->onnx->trtexec
# (optional) Exporting a Model from PyTorch to ONNX and Running it using ONNX Runtime
import torchvision
import torch,os
from models.squeezenet_caffe import squeezenet_caffe
batch_size = 1 # just a random number
current_dir=os.path.dirname(os.path.abspath(__file__)) # 获取当前路径
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = squeezenet_caffe().cuda()
model_path='/data/cch/pytorch-cifar100-master/checkpoint/squeezenet_caffe/Monday_04_September_2023_11h_48m_33s/squeezenet_caffe-297-best.pth' # cloth
state_dict = torch.load(model_path, map_location=device)
print(1)
# mew_state_dict = OrderedDict()
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in state_dict.items() if (k in model_dict and 'fc' not in k)}
model_dict.update(pretrained_dict)
print(2)
model.load_state_dict(model_dict)
model.eval()
print(3)
# output = model(data)
# Input to the model
x = torch.randn(batch_size, 3, 100, 100, requires_grad=True)
x = x.cuda()
torch_out = model(x)
# Export the model
torch.onnx.export(model, # model being run
x, # model input (or a tuple for multiple inputs)
"/data/cch/pytorch-cifar100-master/checkpoint/squeezenet_caffe/Monday_04_September_2023_11h_48m_33s/squeezenet_caffe-297-best.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
opset_version=10, # the ONNX version to export the model to
do_constant_folding=True, # whether to execute constant folding for optimization
input_names = ['input'], # the model's input names
output_names = ['output'], # the model's output names
dynamic_axes={'input' : {0 : 'batch_size'}, # variable length axes
'output' : {0 : 'batch_size'}})
只需要修改一下输入输出的路径和输入的size即可。
然后是onnx转trt，这里需要自己先安装搭建好tensorrt的环境（环境搭建可能会有点复杂需要编译，有时间单独出一个详细的搭建过程），然后在tensorrt工程下的bin目录下运行命令：
./trtexec --onnx=/data/.../best.onnx --saveEngine=/data.../best.trt --workspace=6000
TensorRT可以提供workspace作为每层网络执行时的临时存储空间，该空间是共享的以减少显存占用（单位是M）。具体的原理可以参考这篇。
前向推理
代码如下：
# 动态推理
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import torchvision.transforms as transforms
from PIL import Image
def load_engine(engine_path):
# TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # INFO
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
print('---')
print(trt.Runtime(TRT_LOGGER))
print('---')
with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
# 2. 读取数据，数据处理为可以和网络结构输入对应起来的的shape，数据可增加预处理
def get_test_transform():
return transforms.Compose([
transforms.Resize([100, 100]),
transforms.ToTensor(),
# transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
transforms.Normalize(mean=[0.4796262, 0.4549252, 0.43396652], std=[0.27888104, 0.28492442, 0.27168077])
])
image = Image.open('/data/.../dog.jpg')
image = get_test_transform()(image)
image = image.unsqueeze_(0) # -> NCHW, 1,3,224,224
print("input img mean {} and std {}".format(image.mean(), image.std()))
image = np.array(image)
path = '/data/.../squeezenet_caffe-297-best.trt'
# 1. 建立模型，构建上下文管理器
engine = load_engine(path)
print(engine)
context = engine.create_execution_context()
context.active_optimization_profile = 0
# 3.分配内存空间，并进行数据cpu到gpu的拷贝
# 动态尺寸，每次都要set一下模型输入的shape，0代表的就是输入，输出根据具体的网络结构而定，可以是0,1,2,3...其中的某个头。
context.set_binding_shape(0, image.shape)
d_input = cuda.mem_alloc(image.nbytes) # 分配输入的内存。
output_shape = context.get_binding_shape(1)
buffer = np.empty(output_shape, dtype=np.float32)
d_output = cuda.mem_alloc(buffer.nbytes) # 分配输出内存。
cuda.memcpy_htod(d_input, image)
bindings = [d_input, d_output]
# 4.进行推理，并将结果从gpu拷贝到cpu。
context.execute_v2(bindings) # 可异步和同步
cuda.memcpy_dtoh(buffer, d_output)
output = buffer.reshape(output_shape)
y_pred_binary = np.argmax(output, axis=1)
print(y_pred_binary[0])
