tensorRT 对cuda训练的PyTorch模型加速方法教程，实战总结_使用cuda加速 pycharm代码

一. 流程介绍

本文的tensorRT加速部署的流程主要有四步：

1. 模型导出成 ONNX 格式；
2. 将 ONNX 格式模型输入给 tensor RT ，指定优化参数；
3. 使用 tensor RT 优化得到 tensor RT Engine；
4. 对 tensor RT Engine 进行inference；

二. 模型转换ONNX

这里是上述流程的第一步，直接上PyTorch模型转换的代码，使用的是自带的torch.onnx.export。若是tensorflow模型可以参考这个链接：tensorflow转换为onnx

import torch
import torch.nn
import onnx
from torch.autograd import Variable
import torch.onnx
import torchvision

# 定义一个随机的输入变量，用来遍历模型参数，注意要放入cuda中
dummy_input = Variable(torch.randn(1, 3, 480, 640)).cuda()
# 加载一个PyTorch模型
model = torch.load(r"PyTorch_Model.pt")
# 导出为ONNX格式文件
torch.onnx.export(model, dummy_input, "save_your_onnx_model.onnx", verbose=True)

1
2
3
4
5
6
7
8
9
10
11
12
13
14

转换完后可以通过这个网站上传并查看转换后的onnx模型结构。

三. 从ONNX获取Engine

这一步相对较麻烦，包含了流程的第二步和第三步
1. 查看cuda版本和cudnn版本
要注意与电脑的系统版本（windows/linux/iOS）以及电脑调用tensorRT时要使用的cuda、cudnn版本对应。
①win+R，输入cmd打开终端；

②输入nvcc --version，查看cuda版本；

③进入以下路径C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA，进入你的cuda版本文件夹，我的是v10.2，进入include文件夹，打开cudnn.h文件。

④往下翻一点，找到下面的版本号，如我的就是7.6.5。后来我把cudnn7升级为cudnn8了，tensorrt也换成了8，因为tensorrt7没有python版本，如果有的话请告知。不过这都不重要hhh

2. 下载并配置tensorRT
下载的tensorRT到nvidia官网下载，我的电脑的cuda是10.2.89，cudnn版本是8.x，网上搜到对应的tensorRT版本就是8.x。下载好了之后就需要配置环境变量。在系统环境变量中添加：
你的tensorRT下载路径\TensorRT-8.2.1.8.Windows10.x86_64.cuda-10.2.cudnn8.2\TensorRT-8.2.1.8\bin
你的tensorRT下载路径\TensorRT-8.2.1.8.Windows10.x86_64.cuda-10.2.cudnn8.2\TensorRT-8.2.1.8\lib
3. 使用tensorRT转换成Engine
win+R，输入cmd打开终端，输入trtexec --onnx=model.onnx --saveEngine=xxx.trt，回车，稍等一会，就可以将模型转换成trt文件啦。
4. 转换的时候记得添加优化参数，这里举个例子
trtexec --onnx=model.onnx --saveEngine=xxx.trt --fp16
①--fp16
通常情况下，fp16精度就可以把模型的遍历速度提高一倍，并损失非常少的精度，若使用int8，精度可能会大大降低，同时速度不一定能比fp16更快一倍。
②其他优化参数可以自行网上搜索添加。

######### 在转换过程中，我遇到了报错如下 ##########

①由于找不到nvparsers.dll，无法继续执行代码。重新安装程序可能会解决此问题。

解决方法：
将tensorrt8/lib添加系统环境变量后，就没出现这个报错。
如果还是有错，就把dll复制一份到tensorrt8/bin中试试。

②找不到cudnn64_8.dll，无法继续执行代码。
解决方法 ：重新检查tensorRT版本下载的对不对，cuda和cudnn的版本一定都要校对。

四. 调用Engine加速

C++调用加速，直接贴代码。

//读取 engine 文件
bool Widget::initial_the_trtfile(const std::string& enginePath){

    char* trtModelStream;
    data = new float[1 * 3 * INPUT_H * INPUT_W];
    output = new float[1 * OUTPUT_SIZE];

    std::ifstream file(enginePath, std::ios::binary);
    if (file.good()) {
        file.seekg(0, file.end);
        size = file.tellg();
        file.seekg(0, file.beg);
        trtModelStream = new char[size];
        assert(trtModelStream);
        file.read(trtModelStream, size);
        file.close();
    }

    runtime = nvinfer1::createInferRuntime(gLogger);
    assert(runtime != nullptr);
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
    assert(engine != nullptr);
    context = engine->createExecutionContext();
    assert(context != nullptr);
    delete[] trtModelStream;

    assert(engine->getNbBindings() == 2);
    // In order to bind the buffers, we need to know the names of the input and output tensors.
    // Note that indices are guaranteed to be less than IEngine::getNbBindings()

    for (int i = 0; i < engine->getNbBindings(); i++){
        nvinfer1::Dims dims = engine->getBindingDimensions(i);
        printf("index %d, dims: (",i);
        for (int d = 0; d < dims.nbDims; d++){
            if (d < dims.nbDims - 1)	printf("%d,", dims.d[d]);
            else	printf("%d", dims.d[d]);
        }	printf(")\n");
    }

    const int inputIndex = 0, outputIndex = 1;
    assert(inputIndex == 0);
    assert(outputIndex == 1);
    // Create GPU buffers on device
    cudaMalloc(&buffers[inputIndex], 1 * 3 * INPUT_H * INPUT_W * sizeof(float));
    cudaMalloc(&buffers[outputIndex], 1 * OUTPUT_SIZE * sizeof(float));
    // Create stream
    cudaStreamCreate(&stream);

    return true;
}

//遍历engine
void Widget::Engine_Inference(nvinfer1::IExecutionContext& context, cudaStream_t& stream, void **buffers, float* input, float* output, int batchSize) {
    // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
    cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream);
    context.enqueue(batchSize, buffers, stream, nullptr);
    cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);
}

//这几个参数可以放到头文件，上面初始化trt文件和遍历的两个函数都要用到。
nvinfer1::IExecutionContext* context;
void* buffers[2];
size_t size;
cudaStream_t stream;

//这个是定义输入输出
data = new float[1 * 3 * INPUT_H * INPUT_W];
output = new float[1 * OUTPUT_SIZE];

//初始化trt文件
bool initial_res = initial_the_trtfile("path to your trt file.");
//遍历的时候调用的是初始化trt文件后得到的context和stream，data就是输入（shape应与你设置输入trt文件的大小一致），output就是输出啦
Widget::Engine_Inference(*context, stream, buffers, data, output, batchsize);
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

完结撒花