pytorch搭建squeezenet网络的整套工程，及其转tensorrt进行cuda加速

本来，前辈们用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])