• 模型部署入门教程（七）：TensorRT 模型构建与推理

    模型部署入门教程继续更新啦！相信经过前几期的学习，大家已经对 ONNX 这一中间表示有了一个比较全面的认识，但是在具体的生产环境中，ONNX 模型常常需要被转换成能被具体推理后端使用的模型格式。本篇教程我们就和大家一起来认识大名鼎鼎的推理后端 TensorRT。

    TensorRT 简介

    TensorRT 是由 NVIDIA 发布的深度学习框架，用于在其硬件上运行深度学习推理。TensorRT 提供量化感知训练和离线量化功能，用户可以选择 INT8 和 FP16 两种优化模式，将深度学习模型应用到不同任务的生产部署，如视频流、语音识别、推荐、欺诈检测、文本生成和自然语言处理。TensorRT 经过高度优化，可在 NVIDIA GPU 上运行， 并且可能是目前在 NVIDIA GPU 运行模型最快的推理引擎。关于 TensorRT 更具体的信息可以访问 TensorRT官网 了解。

    安装 TensorRT

    Windows

    默认在一台有 NVIDIA 显卡的机器上，提前安装好 CUDA 和 CUDNN，登录 NVIDIA 官方网站下载和主机 CUDA 版本适配的 TensorRT 压缩包即可。

    以 CUDA 版本是 10.2 为例，选择适配 CUDA 10.2 的 zip 包，下载完成后，有 conda 虚拟环境的用户可以优先切换到虚拟环境中，然后在 powershell 中执行类似如下的命令安装并测试：

    1. cd \the\path\of\tensorrt\zip\file 
    2. Expand-Archive TensorRT-8.2.5.1.Windows10.x86_64.cuda-10.2.cudnn8.2.zip . 
    3. $env:TENSORRT_DIR = "$pwd\TensorRT-8.2.5.1" 
    4. $env:path = "$env:TENSORRT_DIR\lib;" + $env:path 
    5. pip install $env:TENSORRT_DIR\python\tensorrt-8.2.5.1-cp36-none-win_amd64.whl 
    6. python -c "import tensorrt;print(tensorrt.__version__)"

    上述命令会在安装后检查 TensorRT 版本，如果打印结果是 8.2.5.1，说明安装 Python 包成功了。

    Linux

    和在 Windows 环境下安装类似，默认在一台有 NVIDIA 显卡的机器上，提前安装好 CUDA 和 CUDNN，登录 NVIDIA 官方网站下载和主机 CUDA 版本适配的 TensorRT 压缩包即可。

    以 CUDA 版本是 10.2 为例，选择适配 CUDA 10.2 的 tar 包，然后执行类似如下的命令安装并测试：

    1. cd /the/path/of/tensorrt/tar/gz/file 
    2. tar -zxvf TensorRT-8.2.5.1.linux.x86_64-gnu.cuda-10.2.cudnn8.2.tar.gz 
    3. export TENSORRT_DIR=$(pwd)/TensorRT-8.2.5.1 
    4. export LD_LIBRARY_PATH=$TENSORRT_DIR/lib:$LD_LIBRARY_PATH 
    5. pip install TensorRT-8.2.5.1/python/tensorrt-8.2.5.1-cp37-none-linux_x86_64.whl 
    6. python -c "import tensorrt;print(tensorrt.__version__)"

    如果发现打印结果是 8.2.5.1，说明安装 Python 包成功了。

    模型构建

    我们使用 TensorRT 生成模型主要有两种方式：

    1. 直接通过 TensorRT 的 API 逐层搭建网络；
    2. 将中间表示的模型转换成 TensorRT 的模型，比如将 ONNX 模型转换成 TensorRT 模型。

    接下来，我们将用 Python 和 C++ 语言分别使用这两种方式构建 TensorRT 模型，并将生成的模型进行推理。

    直接构建

    利用 TensorRT 的 API 逐层搭建网络，这一过程类似使用一般的训练框架，如使用 Pytorch 或者TensorFlow 搭建网络。需要注意的是对于权重部分，如卷积或者归一化层，需要将权重内容赋值到 TensorRT 的网络中。本文就不详细展示，只搭建一个对输入做池化的简单网络。

    使用 Python API 构建

    首先是使用 Python API 直接搭建 TensorRT 网络，这种方法主要是利用 tensorrt.Builder 的 create_builder_config 和 create_network 功能，分别构建 config 和 network，前者用于设置网络的最大工作空间等参数，后者就是网络主体，需要对其逐层添加内容。

    此外，需要定义好输入和输出名称，将构建好的网络序列化，保存成本地文件。值得注意的是：如果想要网络接受不同分辨率的输入输出，需要使用 tensorrt.Builder 的 create_optimization_profile 函数，并设置最小、最大的尺寸。

    实现代码如下：

    1. import tensorrt as trt 
    2.  
    3. verbose = True 
    4. IN_NAME = 'input' 
    5. OUT_NAME = 'output' 
    6. IN_H = 224 
    7. IN_W = 224 
    8. BATCH_SIZE = 1 
    9.  
    10. EXPLICIT_BATCH = 1 << (int)( 
    11.     trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) 
    12.  
    13. TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger() 
    14. with trt.Builder(TRT_LOGGER) as builder, builder.create_builder_config( 
    15. ) as config, builder.create_network(EXPLICIT_BATCH) as network: 
    16.     # define network 
    17.     input_tensor = network.add_input( 
    18.         name=IN_NAME, dtype=trt.float32, shape=(BATCH_SIZE, 3, IN_H, IN_W)) 
    19.     pool = network.add_pooling( 
    20.         input=input_tensor, type=trt.PoolingType.MAX, window_size=(2, 2)) 
    21.     pool.stride = (2, 2) 
    22.     pool.get_output(0).name = OUT_NAME 
    23.     network.mark_output(pool.get_output(0)) 
    24.  
    25.     # serialize the model to engine file 
    26.     profile = builder.create_optimization_profile() 
    27.     profile.set_shape_input('input', *[[BATCH_SIZE, 3, IN_H, IN_W]]*3)  
    28.     builder.max_batch_size = 1 
    29.     config.max_workspace_size = 1 << 30 
    30.     engine = builder.build_engine(network, config) 
    31.     with open('model_python_trt.engine', mode='wb') as f: 
    32.         f.write(bytearray(engine.serialize())) 
    33.         print("generating file done!")

    使用 C++ API 构建

    对于想要直接用 C++ 语言构建网络的小伙伴来说，整个流程和上述 Python 的执行过程非常类似，需要注意的点主要有：

    1. nvinfer1:: createInferBuilder 对应 Python 中的 tensorrt.Builder，需要传入 ILogger 类的实例，但是 ILogger 是一个抽象类，需要用户继承该类并实现内部的虚函数。不过此处我们直接使用了 TensorRT 包解压后的 samples 文件夹 ../samples/common/logger.h 文件里的实现 Logger 子类。
    2. 设置 TensorRT 模型的输入尺寸，需要多次调用 IOptimizationProfile 的 setDimensions 方法，比 Python 略繁琐一些。IOptimizationProfile 需要用 createOptimizationProfile 函数，对应 Python 的 create_builder_config 函数。

    实现代码如下：

    1. #include  
    2. #include  
    3.  
    4. #include  
    5. #include <../samples/common/logger.h> 
    6.  
    7. using namespace nvinfer1; 
    8. using namespace sample; 
    9.  
    10. const char* IN_NAME = "input"; 
    11. const char* OUT_NAME = "output"; 
    12. static const int IN_H = 224; 
    13. static const int IN_W = 224; 
    14. static const int BATCH_SIZE = 1; 
    15. static const int EXPLICIT_BATCH = 1 << (int)(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 
    16.  
    17. int main(int argc, char** argv) 
    18. { 
    19.         // Create builder 
    20.         Logger m_logger; 
    21.         IBuilder* builder = createInferBuilder(m_logger); 
    22.         IBuilderConfig* config = builder->createBuilderConfig(); 
    23.  
    24.         // Create model to populate the network 
    25.         INetworkDefinition* network = builder->createNetworkV2(EXPLICIT_BATCH); 
    26.         ITensor* input_tensor = network->addInput(IN_NAME, DataType::kFLOAT, Dims4{ BATCH_SIZE, 3, IN_H, IN_W }); 
    27.         IPoolingLayer* pool = network->addPoolingNd(*input_tensor, PoolingType::kMAX, DimsHW{ 2, 2 }); 
    28.         pool->setStrideNd(DimsHW{ 2, 2 }); 
    29.         pool->getOutput(0)->setName(OUT_NAME); 
    30.         network->markOutput(*pool->getOutput(0)); 
    31.  
    32.         // Build engine 
    33.         IOptimizationProfile* profile = builder->createOptimizationProfile(); 
    34.         profile->setDimensions(IN_NAME, OptProfileSelector::kMIN, Dims4(BATCH_SIZE, 3, IN_H, IN_W)); 
    35.         profile->setDimensions(IN_NAME, OptProfileSelector::kOPT, Dims4(BATCH_SIZE, 3, IN_H, IN_W)); 
    36.         profile->setDimensions(IN_NAME, OptProfileSelector::kMAX, Dims4(BATCH_SIZE, 3, IN_H, IN_W)); 
    37.         config->setMaxWorkspaceSize(1 << 20); 
    38.         ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); 
    39.  
    40.         // Serialize the model to engine file 
    41.         IHostMemory* modelStream{ nullptr }; 
    42.         assert(engine != nullptr); 
    43.         modelStream = engine->serialize(); 
    44.  
    45.         std::ofstream p("model.engine", std::ios::binary); 
    46.         if (!p) { 
    47.                 std::cerr << "could not open output file to save model" << std::endl; 
    48.                 return -1; 
    49.         } 
    50.         p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); 
    51.         std::cout << "generating file done!" << std::endl; 
    52.  
    53.         // Release resources 
    54.         modelStream->destroy(); 
    55.         network->destroy(); 
    56.         engine->destroy(); 
    57.         builder->destroy(); 
    58.         config->destroy(); 
    59.         return 0; 
    60. }

    IR 转换模型

    除了直接通过 TensorRT 的 API 逐层搭建网络并序列化模型，TensorRT 还支持将中间表示的模型（如 ONNX）转换成 TensorRT 模型。

    使用 Python API 转换

    我们首先使用 Pytorch 实现一个和上文一致的模型，即只对输入做一次池化并输出；然后将 Pytorch 模型转换成 ONNX 模型；最后将 ONNX 模型转换成 TensorRT 模型。

    这里主要使用了 TensorRT 的 OnnxParser 功能，它可以将 ONNX 模型解析到 TensorRT 的网络中。最后我们同样可以得到一个 TensorRT 模型，其功能与上述方式实现的模型功能一致。

    实现代码如下：

    1. import torch 
    2. import onnx 
    3. import tensorrt as trt 
    4.  
    5.  
    6. onnx_model = 'model.onnx' 
    7.  
    8. class NaiveModel(torch.nn.Module): 
    9.     def __init__(self): 
    10.         super().__init__() 
    11.         self.pool = torch.nn.MaxPool2d(2, 2) 
    12.  
    13.     def forward(self, x): 
    14.         return self.pool(x) 
    15.  
    16. device = torch.device('cuda:0') 
    17.  
    18. # generate ONNX model 
    19. torch.onnx.export(NaiveModel(), torch.randn(1, 3, 224, 224), onnx_model, input_names=['input'], output_names=['output'], opset_version=11) 
    20. onnx_model = onnx.load(onnx_model) 
    21.  
    22. # create builder and network 
    23. logger = trt.Logger(trt.Logger.ERROR) 
    24. builder = trt.Builder(logger) 
    25. EXPLICIT_BATCH = 1 << (int)( 
    26.     trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) 
    27. network = builder.create_network(EXPLICIT_BATCH) 
    28.  
    29. # parse onnx 
    30. parser = trt.OnnxParser(network, logger) 
    31.  
    32. if not parser.parse(onnx_model.SerializeToString()): 
    33.     error_msgs = '' 
    34.     for error in range(parser.num_errors): 
    35.         error_msgs += f'{parser.get_error(error)}\n' 
    36.     raise RuntimeError(f'Failed to parse onnx, {error_msgs}') 
    37.  
    38. config = builder.create_builder_config() 
    39. config.max_workspace_size = 1<<20 
    40. profile = builder.create_optimization_profile() 
    41.  
    42. profile.set_shape('input', [1,3 ,224 ,224], [1,3,224, 224], [1,3 ,224 ,224]) 
    43. config.add_optimization_profile(profile) 
    44. # create engine 
    45. with torch.cuda.device(device): 
    46.     engine = builder.build_engine(network, config) 
    47.  
    48. with open('model.engine', mode='wb') as f: 
    49.     f.write(bytearray(engine.serialize())) 
    50.     print("generating file done!")

    IR 转换时，如果有多 Batch、多输入、动态 shape 的需求，都可以通过多次调用 set_shape 函数进行设置。set_shape 函数接受的传参分别是：输入节点名称，可接受的最小输入尺寸，最优的输入尺寸，可接受的最大输入尺寸。一般要求这三个尺寸的大小关系为单调递增。

    使用 C++ API 转换

    介绍了如何用 Python 语言将 ONNX 模型转换成 TensorRT 模型后，再介绍下如何用 C++ 将 ONNX 模型转换成 TensorRT 模型。这里通过 NvOnnxParser，我们可以将上一小节转换时得到的 ONNX 文件直接解析到网络中。

    实现代码如下：

    1. #include  
    2. #include  
    3.  
    4. #include  
    5. #include  
    6. #include <../samples/common/logger.h> 
    7.  
    8. using namespace nvinfer1; 
    9. using namespace nvonnxparser; 
    10. using namespace sample; 
    11.  
    12. int main(int argc, char** argv) 
    13. { 
    14.         // Create builder 
    15.         Logger m_logger; 
    16.         IBuilder* builder = createInferBuilder(m_logger); 
    17.         const auto explicitBatch = 1U << static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 
    18.         IBuilderConfig* config = builder->createBuilderConfig(); 
    19.  
    20.         // Create model to populate the network 
    21.         INetworkDefinition* network = builder->createNetworkV2(explicitBatch); 
    22.  
    23.         // Parse ONNX file 
    24.         IParser* parser = nvonnxparser::createParser(*network, m_logger); 
    25.         bool parser_status = parser->parseFromFile("model.onnx", static_cast<int>(ILogger::Severity::kWARNING)); 
    26.  
    27.         // Get the name of network input 
    28.         Dims dim = network->getInput(0)->getDimensions(); 
    29.         if (dim.d[0] == -1)  // -1 means it is a dynamic model 
    30.         { 
    31.                 const char* name = network->getInput(0)->getName(); 
    32.                 IOptimizationProfile* profile = builder->createOptimizationProfile(); 
    33.                 profile->setDimensions(name, OptProfileSelector::kMIN, Dims4(1, dim.d[1], dim.d[2], dim.d[3])); 
    34.                 profile->setDimensions(name, OptProfileSelector::kOPT, Dims4(1, dim.d[1], dim.d[2], dim.d[3])); 
    35.                 profile->setDimensions(name, OptProfileSelector::kMAX, Dims4(1, dim.d[1], dim.d[2], dim.d[3])); 
    36.                 config->addOptimizationProfile(profile); 
    37.         } 
    38.  
    39.  
    40.         // Build engine 
    41.         config->setMaxWorkspaceSize(1 << 20); 
    42.         ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); 
    43.  
    44.         // Serialize the model to engine file 
    45.         IHostMemory* modelStream{ nullptr }; 
    46.         assert(engine != nullptr); 
    47.         modelStream = engine->serialize(); 
    48.  
    49.         std::ofstream p("model.engine", std::ios::binary); 
    50.         if (!p) { 
    51.                 std::cerr << "could not open output file to save model" << std::endl; 
    52.                 return -1; 
    53.         } 
    54.         p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); 
    55.         std::cout << "generate file success!" << std::endl; 
    56.  
    57.         // Release resources 
    58.         modelStream->destroy(); 
    59.         network->destroy(); 
    60.         engine->destroy(); 
    61.         builder->destroy(); 
    62.         config->destroy(); 
    63.         return 0; 
    64. } 
    65.

    模型推理

    前面，我们使用了两种构建 TensorRT 模型的方式，分别用 Python 和 C++ 两种语言共生成了四个 TensorRT 模型，这四个模型的功能理论上是完全一致的。

    接下来，我们将分别使用 Python 和 C++ 两种语言对生成的 TensorRT 模型进行推理。

    使用 Python API 推理

    首先是使用 Python API 推理 TensorRT 模型，这里部分代码参考了 MMDeploy。运行下面代码，可以发现输入一个 1x3x224x224 的张量，输出一个 1x3x112x112 的张量，完全符合我们对输入池化后结果的预期。

    1. from typing import Union, Optional, Sequence,Dict,Any 
    2.  
    3. import torch 
    4. import tensorrt as trt 
    5.  
    6. class TRTWrapper(torch.nn.Module): 
    7.     def __init__(self,engine: Union[str, trt.ICudaEngine], 
    8.                  output_names: Optional[Sequence[str]] = None) -> None: 
    9.         super().__init__() 
    10.         self.engine = engine 
    11.         if isinstance(self.engine, str): 
    12.             with trt.Logger() as logger, trt.Runtime(logger) as runtime: 
    13.                 with open(self.engine, mode='rb') as f: 
    14.                     engine_bytes = f.read() 
    15.                 self.engine = runtime.deserialize_cuda_engine(engine_bytes) 
    16.         self.context = self.engine.create_execution_context() 
    17.         names = [_ for _ in self.engine] 
    18.         input_names = list(filter(self.engine.binding_is_input, names)) 
    19.         self._input_names = input_names 
    20.         self._output_names = output_names 
    21.  
    22.         if self._output_names is None: 
    23.             output_names = list(set(names) - set(input_names)) 
    24.             self._output_names = output_names 
    25.  
    26.     def forward(self, inputs: Dict[str, torch.Tensor]): 
    27.         assert self._input_names is not None 
    28.         assert self._output_names is not None 
    29.         bindings = [None] * (len(self._input_names) + len(self._output_names)) 
    30.         profile_id = 0 
    31.         for input_name, input_tensor in inputs.items(): 
    32.             # check if input shape is valid 
    33.             profile = self.engine.get_profile_shape(profile_id, input_name) 
    34.             assert input_tensor.dim() == len( 
    35.                 profile[0]), 'Input dim is different from engine profile.' 
    36.             for s_min, s_input, s_max in zip(profile[0], input_tensor.shape, 
    37.                                              profile[2]): 
    38.                 assert s_min <= s_input <= s_max, \ 
    39.                     'Input shape should be between ' \ 
    40.                     + f'{profile[0]} and {profile[2]}' \ 
    41.                     + f' but get {tuple(input_tensor.shape)}.' 
    42.             idx = self.engine.get_binding_index(input_name) 
    43.  
    44.             # All input tensors must be gpu variables 
    45.             assert 'cuda' in input_tensor.device.type 
    46.             input_tensor = input_tensor.contiguous() 
    47.             if input_tensor.dtype == torch.long: 
    48.                 input_tensor = input_tensor.int() 
    49.             self.context.set_binding_shape(idx, tuple(input_tensor.shape)) 
    50.             bindings[idx] = input_tensor.contiguous().data_ptr() 
    51.  
    52.         # create output tensors 
    53.         outputs = {} 
    54.         for output_name in self._output_names: 
    55.             idx = self.engine.get_binding_index(output_name) 
    56.             dtype = torch.float32 
    57.             shape = tuple(self.context.get_binding_shape(idx)) 
    58.  
    59.             device = torch.device('cuda') 
    60.             output = torch.empty(size=shape, dtype=dtype, device=device) 
    61.             outputs[output_name] = output 
    62.             bindings[idx] = output.data_ptr() 
    63.         self.context.execute_async_v2(bindings, 
    64.                                       torch.cuda.current_stream().cuda_stream) 
    65.         return outputs 
    66.  
    67. model = TRTWrapper('model.engine', ['output']) 
    68. output = model(dict(input = torch.randn(1, 3, 224, 224).cuda())) 
    69. print(output)

    使用 C++ API 推理

    最后，在很多实际生产环境中，我们都会使用 C++ 语言完成具体的任务，以达到更加高效的代码运行效果，另外 TensoRT 的用户一般也都更看重其在 C++ 下的使用，所以我们也用 C++ 语言实现一遍模型推理，这也可以和用 Python API 推理模型做一个对比。

    实现代码如下：

    1. #include  
    2. #include  
    3.  
    4. #include  
    5. #include <../samples/common/logger.h> 
    6.  
    7. #define CHECK(status) \ 
    8.     do\ 
    9.     {\ 
    10.         auto ret = (status);\ 
    11.         if (ret != 0)\ 
    12.         {\ 
    13.             std::cerr << "Cuda failure: " << ret << std::endl;\ 
    14.             abort();\ 
    15.         }\ 
    16.     } while (0) 
    17.  
    18. using namespace nvinfer1; 
    19. using namespace sample; 
    20.  
    21. const char* IN_NAME = "input"; 
    22. const char* OUT_NAME = "output"; 
    23. static const int IN_H = 224; 
    24. static const int IN_W = 224; 
    25. static const int BATCH_SIZE = 1; 
    26. static const int EXPLICIT_BATCH = 1 << (int)(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); 
    27.  
    28.  
    29. void doInference(IExecutionContext& context, float* input, float* output, int batchSize) 
    30. { 
    31.         const ICudaEngine& engine = context.getEngine(); 
    32.  
    33.         // Pointers to input and output device buffers to pass to engine. 
    34.         // Engine requires exactly IEngine::getNbBindings() number of buffers. 
    35.         assert(engine.getNbBindings() == 2); 
    36.         void* buffers[2]; 
    37.  
    38.         // In order to bind the buffers, we need to know the names of the input and output tensors. 
    39.         // Note that indices are guaranteed to be less than IEngine::getNbBindings() 
    40.         const int inputIndex = engine.getBindingIndex(IN_NAME); 
    41.         const int outputIndex = engine.getBindingIndex(OUT_NAME); 
    42.  
    43.         // Create GPU buffers on device 
    44.         CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * IN_H * IN_W * sizeof(float))); 
    45.         CHECK(cudaMalloc(&buffers[outputIndex], batchSize * 3 * IN_H * IN_W /4 * sizeof(float))); 
    46.  
    47.         // Create stream 
    48.         cudaStream_t stream; 
    49.         CHECK(cudaStreamCreate(&stream)); 
    50.  
    51.         // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host 
    52.         CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * IN_H * IN_W * sizeof(float), cudaMemcpyHostToDevice, stream)); 
    53.         context.enqueue(batchSize, buffers, stream, nullptr); 
    54.         CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * 3 * IN_H * IN_W / 4 * sizeof(float), cudaMemcpyDeviceToHost, stream)); 
    55.         cudaStreamSynchronize(stream); 
    56.  
    57.         // Release stream and buffers 
    58.         cudaStreamDestroy(stream); 
    59.         CHECK(cudaFree(buffers[inputIndex])); 
    60.         CHECK(cudaFree(buffers[outputIndex])); 
    61. } 
    62.  
    63. int main(int argc, char** argv) 
    64. { 
    65.         // create a model using the API directly and serialize it to a stream 
    66.         char *trtModelStream{ nullptr }; 
    67.         size_t size{ 0 }; 
    68.  
    69.         std::ifstream file("model.engine", std::ios::binary); 
    70.         if (file.good()) { 
    71.                 file.seekg(0, file.end); 
    72.                 size = file.tellg(); 
    73.                 file.seekg(0, file.beg); 
    74.                 trtModelStream = new char[size]; 
    75.                 assert(trtModelStream); 
    76.                 file.read(trtModelStream, size); 
    77.                 file.close(); 
    78.         } 
    79.  
    80.         Logger m_logger; 
    81.         IRuntime* runtime = createInferRuntime(m_logger); 
    82.         assert(runtime != nullptr); 
    83.         ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); 
    84.         assert(engine != nullptr); 
    85.         IExecutionContext* context = engine->createExecutionContext(); 
    86.         assert(context != nullptr); 
    87.  
    88.         // generate input data 
    89.         float data[BATCH_SIZE * 3 * IN_H * IN_W]; 
    90.         for (int i = 0; i < BATCH_SIZE * 3 * IN_H * IN_W; i++) 
    91.                 data[i] = 1; 
    92.  
    93.         // Run inference 
    94.         float prob[BATCH_SIZE * 3 * IN_H * IN_W /4]; 
    95.         doInference(*context, data, prob, BATCH_SIZE); 
    96.  
    97.         // Destroy the engine 
    98.         context->destroy(); 
    99.         engine->destroy(); 
    100.         runtime->destroy(); 
    101.         return 0; 
    102. }

    总结

    通过本文的学习，我们掌握了两种构建 TensorRT 模型的方式：直接通过 TensorRT 的 API 逐层搭建网络；将中间表示的模型转换成 TensorRT 的模型。不仅如此，我们还分别用 C++ 和 Python 两种语言完成了 TensorRT 模型的构建及推理，相信大家都有所收获！在下一篇文章中，我们将和大家一起学习何添加 TensorRT 自定义算子，敬请期待哦~

    FAQ

    • Q：运行代码时报错：Could not find: cudnn64_8.dll. Is it on your PATH?
    • A：首先检查下自己的环境变量中是否包含 cudnn64_8.dll 所在的路径，若发现 cudnn 的路径在 C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v10.2\\bin 中，但是里面只有 cudnn64_7.dll，解决方法是去 NVIDIA 官网下载 cuDNN zip 包，解压后，复制其中的 cudnn64_8.dll 到 CUDA Toolkit 的 bin 目录下。这时也可以复制一份 cudnn64_7.dll，然后将复制的那份改名成 cudnn64_8.dll，同样可以解决这个问题。

    参考

    GitHub - wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT networ

    GitHub - NVIDIA/TensorRT: TensorRT is a C++ library for high performance inference on NVIDIA GPUs an

    系列传送门

    OpenMMLab：模型部署入门教程（一）：模型部署简介172 赞同 · 22 评论文章213 赞同 · 24 评论文章241 赞同 · 31 评论文章241 赞同 · 31 评论文章正在上传…重新上传取消

    OpenMMLab：模型部署入门教程（二）：解决模型部署中的难题77 赞同 · 22 评论文章89 赞同 · 27 评论文章96 赞同 · 27 评论文章96 赞同 · 27 评论文章正在上传…重新上传取消

    OpenMMLab：模型部署入门教程（三）：PyTorch 转 ONNX 详解131 赞同 · 25 评论文章146 赞同 · 32 评论文章153 赞同 · 34 评论文章153 赞同 · 34 评论文章正在上传…重新上传取消

    OpenMMLab：模型部署入门教程（四）：在 PyTorch 中支持更多 ONNX 算子68 赞同 · 19 评论文章80 赞同 · 28 评论文章86 赞同 · 32 评论文章86 赞同 · 32 评论文章正在上传…重新上传取消

    OpenMMLab：模型部署入门教程（五）：ONNX 模型的修改与调试86 赞同 · 4 评论文章115 赞同 · 10 评论文章122 赞同 · 16 评论文章122 赞同 · 16 评论文章正在上传…重新上传取消

    OpenMMLab：模型部署入门教程（六）：实现 PyTorch-ONNX 精度对齐工具91 赞同 · 7 评论文章正在上传…重新上传取消

    OpenMMLab：TorchScript 解读（一）：初识 TorchScript83 赞同 · 10 评论文章106 赞同 · 15 评论文章118 赞同 · 15 评论文章118 赞同 · 15 评论文章正在上传…重新上传取消

    OpenMMLab：TorchScript 解读（二）：Torch jit tracer 实现解析39 赞同 · 2 评论文章46 赞同 · 2 评论文章46 赞同 · 2 评论文章46 赞同 · 2 评论文章正在上传…重新上传取消

    OpenMMLab：TorchScript 解读（三）：jit 中的 subgraph rewriter24 赞同 · 0 评论文章29 赞同 · 0 评论文章32 赞同 · 0 评论文章32 赞同 · 0 评论文章正在上传…重新上传取消

    OpenMMLab：TorchScript 解读（四）：Torch jit 中的别名分析7 赞同 · 0 评论文章14 赞同 · 0 评论文章16 赞同 · 0 评论文章16 赞同 · 0 评论文章正在上传…重新上传取消

  • 相关阅读:
    Mysql优化---锁机制、行锁及表锁
    LuatOS-SOC接口文档(air780E)--mcu - 封装mcu一些特殊操作
    排序算法时间复杂
    【D3.js】1.18-给 D3 标签添加样式
    【AICFD案例操作】汽车外气动分析
    线框图软件：Balsamiq Wireframes mac中文介绍
    数据库调优：Mysql索引对group by 排序的影响
    学点Java打小工_Day4_数组_冒泡排序
    Mysql提取表字段中的字符串
    LuatOS-SOC接口文档(air780E)--dac - 数模转换
  • 原文地址：https://blog.csdn.net/qq_39967751/article/details/126061511