前面写了一系列YOLO系列的剪枝以及libtorch推理的文章，如何将模型进行部署这才是一个难点，将学术成果转为实际应用，因此后续将不定期写一些有关tensorRT学习的文章，本文也是参考以下链接的教程，相信使用tensorrt推理yolov5大家都用过了。

https://github.com/wang-xinyu/pytorchx

本文将以lenet为例进行tensorRT的推理。

环境说明：

pytorch 1.7

tensorRT: 8.2.4.2

cuda:10.2

windows 10

NVIDIA GTX 1650

首先clone代码到本地。

git clone https://github.com/wang-xinyu/pytorchx

cd pytorchx/lenet

运行项目下的lenet5.py，将在当前路径下保存lenet5.pth文件，这个就是我们经常看到pytorch下保存的权重文件，不过这里需要注意一点的是，这里的pth中保存了权值以及完整的网络模型，从torch.save(mode,..)就可以看出来，而不是用的torch.save(model.state_dict())保存的。

torch转wts

接下是将pth转wts文件。运行项目下的inference.py，完整代码如下：

由于pth文件中已经含有了完整的网络结构，因此使用torch.load的时候就直接将权值已经加载进了网络中。

网络结构参数量的写入

下面这一行是将网络的长度或者说层数先获取并且写入。

f.write("{}\n".format(len(net.state_dict().keys())))

下面这一行代码是将每层的value值铺平，比如第一层的卷积为Conv2d(1,6,k=5,s=1,p=0)，那么对应的value的shape就为(6,1,5,5)，铺平以后为shape为【150】。

vr = v.reshape(-1).cpu().numpy()

从 f.write('{} {}\n'.format(k, len(vr))) 可以判断出写入的时候为当前层所对应value的shape,形式如下：

conv1.weight 150

conv1.bias 6

可以看到会将每层的weight 与bias的shape都写入。

网络结构参数的写入

在看下面这一段代码，由于我们上面拉平后的张量转为了array形式，因此我们通过for循环获取每个weight和bias的值并写入。struck.pack将Python的值根据格式符，转换为字符串(因为Python中没有字节(Byte)类型，可以把这里的字符串理解为字节流，或字节数组)。

hex()函数：用于获取给定数字的十六进制值，它接受一个数字并返回一个十六进制值(字符串格式)

for vv in vr:

   f.write(" ")

   f.write(struct.pack(">f", float(vv)).hex())

写入的形式如下：这将会包含每层的名字(key值)，对应层的参数长度以及每个每个参数的16进制。

conv1.weight 150

be2005c5 bd622f78 be22672a 3e024f0b 3e3e63c1 bd878a9a 3e166093 3dafc9d6 3d666aa4 3cdce100 3e24afad 3e3a9d8d be37de00 bdf5162a 3de03ed2 be4c3bd5 3c1910b0 3df8e38a be102f73 be0b8c3e bd14298c bdbdaca7 3d6b928c 3e164421 bd905b6a bdd32847 3e0c9061 3e3425d7 3e4cc113 3e165f7f 3e2ccce1 bb93b640 3e124f37 3e197a65 bd5b3660 3c4b21b0 3d799e04 be1b1910 3dcad0d2 3e1e2493 3d934aae 3d71c834 bcbc1610 bdeb1bc0 bda3570a 3c91f710 bc76a660 3dc6d17e 3e3d92d3 3e2bf3b9 bce63bc0 3d988152 3c29ded0 bd20866c 3dcd6b32 bdf5d6b3 3d576284 bd9ce097 3e2cfc2b be2eb9af 3c876630 3d9cb1aa 3e10120b 3e1d5497 3dcaabf2 bdc35ab0 3d862dd2 3da98416 be2e4e93 be2e2508 bd8e965a be1bd880 bde0835a 3d9505ee bd15c3e0 3d5a745c 3cbbeb30 3e00388b 3daf33f6 bd847c54 bd39f1d8 bcb15e10 3d527f04 bd80f526 3dffc606 3d84317a 3e3c7153 3d67316c 3de65f5a 3cc4d5e8 bd810c8a bceebf30 be389112 bb0324c0 be341860 be11c97d 3b977360 3e2d2c85 3ca4df90 3e22d299 be1afccb bdf9a6ca bdfe5253 bdeb282a bd94143c 3cd48218 be0b7295 be2cd3aa 3dc5acb6 be39cdf3 bd193b54 be3d2eb0 3e3368ef be360200 bd5c2720 be15b16a 3c0a04e0 3d584104 3d990402 3e184ff5 be390608 be16ab2b bd7a848c bc2c0ba0 3e002d1d 3e0df3db 3de5fe3a 3b2f6280 3e1d28c3 bbe82aa0 bd7f8800 be24cf9e 3e20421f bdefa797 3e378c5b bda28a13 3da8e526 3ddda312 3e375b01 be4bfad7 bc00ed30 3dedcb6e be3188b7 bb0d1800 be00b4fa 3cdfc230 3e2253ad 3e2f56a7 bd33ec2c bdaa9750

conv1.bias 6

be0c3f24 bc742560 3dbe9fee 3d99b8b2 bcac1f80 3df6c97e

import torch
from torch import nn
from lenet5 import Lenet5
import os
import struct
def main():
    print('cuda device count: ', torch.cuda.device_count())
    net = torch.load('lenet5.pth')
    net = net.to('cuda:0')
    net.eval()
    #print('model: ', net)
    #print('state dict: ', net.state_dict()['conv1.weight'])
    tmp = torch.ones(1, 1, 32, 32).to('cuda:0')
    #print('input: ', tmp)
    out = net(tmp)
    print('lenet out:', out)
    f = open("lenet5.wts", 'w')
    f.write("{}\n".format(len(net.state_dict().keys())))
    for k,v in net.state_dict().items():
        #print('key: ', k)
        #print('value: ', v.shape)
        vr = v.reshape(-1).cpu().numpy()
        f.write("{} {}".format(k, len(vr)))
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())
        f.write("\n")
if __name__ == '__main__':
    main()

运行上面的代码以后会将整个的网络进行写入【这里的每层的参数都是拉平写入的，即变成一行】。

在网络最终的输出的grad_fn也可以看到最终的输出流操作为Softmax：

lenet out: tensor([[0.0950, 0.0998, 0.1101, 0.0975, 0.0966, 0.1097, 0.0948, 0.1056, 0.0992,

        0.0917]], device='cuda:0', grad_fn=)

通过上面的步骤就完成了torch转wts。

tensoRT运行lenet5

这里需要提前安装好tensorRT.

tensorRT的安装与配置可以参考：

tensorrt部署YOLOv5模型记录【附代码，支持视频检测】_爱吃肉的鹏的博客-CSDN博客_tensorrt yolo

wts转engine

如果你是linux系统：

git clone https://github.com/wang-xinyu/tensorrtx

cd tensorrtx/lenet

cp [PATH-OF-pytorchx]/pytorchx/lenet/lenet5.wts .

mkdir build

cd build

cmake ..

make

make成功以后需要将wts转为engine文件，使其变为序列化文件。

./lenet -s

这个过程会比较长，耐心等待！

等成功生成engine文件过，就可以进行反序列化的推理了。

./lenet -d

如果你是windows系统：

git clone https://github.com/wang-xinyu/tensorrtx

cd tensorrtx/lenet

在lenet中新建build文件夹。

将上面生成的wts文件复制到lenet下。

我这里是在windows环境下进行的，所以CMakeList.txt修改如下：

cmake_minimum_required(VERSION 2.6)

project(lenet)

add_definitions(-std=c++11)

set(TARGET_NAME "lenet")

set(TRT_DIR "F:\\TensorRT-8.2.4.2") # tensorrt路径

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)

set(CMAKE_CXX_STANDARD 11)

set(CMAKE_BUILD_TYPE Debug)

include_directories(${PROJECT_SOURCE_DIR}/include)

# include and link dirs of cuda and tensorrt, you need adapt them if yours are different

# cuda 设置cuda include和lib路径

enable_language(CUDA)  # add this line, then no need to setup cuda path in vs

include_directories(C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.2/include)

link_directories(C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v10.2/lib/x64)

# tensorrt

include_directories(${TRT_DIR}/include)

link_directories(${TRT_DIR}/lib)

FILE(GLOB SRC_FILES ${PROJECT_SOURCE_DIR}/lenet.cpp ${PROJECT_SOURCE_DIR}/include/*.h)

add_executable(${TARGET_NAME} ${SRC_FILES})

target_link_libraries(${TARGET_NAME} nvinfer)

target_link_libraries(${TARGET_NAME} cudart)

add_definitions(-O2 -pthread)

打开CMake软件(这个软件需要自己下载好)。

输入lenet路径和build的路径。

依次点：configure【会跳出一个窗口进行配置即可，我这里是VS2017】，然后点击Generate，在点击Open Project，此时会自动打开VS编译器。

打开VS以后修改如下：

然后右键你的项目，点击重新生成即可。

2>C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include\cuda_runtime_api.h(9348): warning C4819: 该文件包含不能在当前代码页(936)中表示的字符。请将该文件保存为 Unicode 格式以防止数据丢失

2>F:\TensorRT-8.2.4.2\include\NvInfer.h(2689): warning C4819: 该文件包含不能在当前代码页(936)中表示的字符。请将该文件保存为 Unicode 格式以防止数据丢失

2>F:\TensorRT-8.2.4.2\include\NvInfer.h(4402): warning C4819: 该文件包含不能在当前代码页(936)中表示的字符。请将该文件保存为 Unicode 格式以防止数据丢失

2>F:\TensorRT-8.2.4.2\include\NvInfer.h(5299): warning C4819: 该文件包含不能在当前代码页(936)中表示的字符。请将该文件保存为 Unicode 格式以防止数据丢失

2>C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include\cuda_runtime_api.h(9348): warning C4819: 该文件包含不能在当前代码页(936)中表示的字符。请将该文件保存为 Unicode 格式以防止数据丢失

2>lenet.vcxproj -> E:\tensorrtx\lenet\build\Release\lenet.exe

2>已完成生成项目“lenet.vcxproj”的操作。

========== 全部重新生成: 成功 2 个，失败 0 个，跳过 0 个 ==========

打开cmd，输入lenet.exe -s

生成engine文件【该过程比较长，耐心等待】。如果在cmd中生成engine一直没反应，那么可以进入代码，手动传参试试。

Loading weights: ../lenet5.wts

[10/01/2022-15:16:00] [W] [TRT] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0

[10/01/2022-15:16:09] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:16:09] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:16:10] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:16:10] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:16:10] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:16:11] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:05] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:05] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:05] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:05] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] Try increasing the workspace size to 4194304 bytes to get better performance.

[10/01/2022-15:19:07] [W] [TRT] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0

此时会生成engine文件。

在cmd中输入 lenet.exe -d，可以看到以下输出：

E:\tensorrtx\lenet\build\Release>lenet.exe -d

[10/01/2022-15:21:41] [W] [TRT] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0

[10/01/2022-15:21:42] [W] [TRT] TensorRT was linked against cuDNN 8.2.1 but loaded cuDNN 8.2.0

Output:

0.0949623, 0.0998472, 0.110072, 0.0975036, 0.0965564, 0.109736, 0.0947979, 0.105618, 0.099228, 0.0916792,

onnx转engine

上面的engine是通过将wts转engine，这种方法是可行的，但也可以发现一个问题，就是需要用C++再重新写一遍网络结构，然后将wts中的权重传入构建的网络中，如果网络很简单是可以的，但如果网络结构复杂就肯定会带来很多问题。而tensorRT自带一种可以将onnx转engine的方法，该方法不需要再用C++重写一遍网络结构。

首先需要将torch文件转为onnx，这个网上有很多代码。

torch转onnx：

import torch
from lenet5 import Lenet5
def Torch2Onnx(model_path, onnx_path):
    device = torch.device('cuda:0')
    model = torch.load(model_path, map_location=device)
    test_input = torch.randn(1, 1, 32, 32).to(device)
    input_names = ['data']
    output_names = ['prob']
    torch.onnx.export(
        model,
        test_input,
        onnx_path,
        verbose=True,
        opset_version=11,
        input_names = input_names,
        output_names= output_names
    )
    print("模型转换完成")
if __name__ == "__main__":
    model_path = './lenet5.pth'
    Torch2Onnx(model_path, 'lenet5.onnx')

onnx转engine：

找到你的TensorRT文件，找到bin目录下有个trtexec.exe ,并将你的onnx放入该路径下。输入命令：

$ ./trtexec --onnx=lenet5.onnx --saveEngine=lenet5.engine --workspace=1024

[10/08/2022-11:33:57] [I] D2H Latency: min = 0.00769043 ms, max = 0.723877 ms, mean = 0.0172087 ms, median = 0.0166016 ms, percentile(99%) = 0.052002 ms
[10/08/2022-11:33:57] [I] Total Host Walltime: 3.00059 s
[10/08/2022-11:33:57] [I] Total GPU Compute Time: 0.944105 s
[10/08/2022-11:33:57] [I] Explanations of the performance metrics are printed in the verbose logs.
[10/08/2022-11:33:57] [I]
&&&& PASSED TensorRT.trtexec [TensorRT v8204] # F:\TensorRT-8.2.4.2\bin\trtexec.exe --onnx=lenet5.onnx --saveEngine=lenet5.engine --workspace=1024

此时就会自动生成engine文件，然后我们再将该文件复制到C++工程中即可正常推理。

【注意：这里有点需要注意一下，转onnx的时候需要输入input_name和output_name，这个两个名字必须和C++中的Input和ouput名字一样，不然推理失败】

所遇问题：

错误 C2694 “void Logger::log(nvinfer1::ILogger::Severity,const char *)”: 重写虚函数的限制性异常规范比基类虚成员函数

解决：由于我这里是tensorRT 8.x版本，因此找到一下代码：

void Logger::log(Severity severity, const char* msg) override

改为：

void Logger::log(Severity severity, const char* msg) noexcept override