• tensorrt的安装和使用

    安装

    提前安装好 CUDA 和 CUDNN，登录 NVIDIA 官方网站下载和主机 CUDA 版本适配的 TensorRT 压缩包即可。

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

    1. #安装c++版本
    2. cd /the/path/of/tensorrt/tar/gz/file 
    3. tar -zxvf TensorRT-8.2.5.1.linux.x86_64-gnu.cuda-10.2.cudnn8.2.tar.gz 
    4. export TENSORRT_DIR=$(pwd)/TensorRT-8.2.5.1 
    5. export LD_LIBRARY_PATH=$TENSORRT_DIR/lib:$LD_LIBRARY_PATH 
    6.  
    7. #安装python版本
    8. pip install TensorRT-8.2.5.1/python/tensorrt-8.2.5.1-cp37-none-linux_x86_64.whl 
    9. python -c "import tensorrt;print(tensorrt.__version__)" #打印8.2.5.1，则说明安装成功
    10. #pip install --no-cache-dir --extra-index-url https://pypi.nvidia.com tensorrt==9.2.0.post12.dev也可以

    构建trt模型

    手动搭建

    使用python接口

    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++接口

    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. }

    onnx模型转换

    trtexec

    失败尝试：下载tensorrt的官方镜像，试图使用镜像中的trtexec，但是在驱动版本为515.105的A800机器上，不论哪个版本的镜像，都会报cuda版本跟驱动版本不匹配的错误

    1. docker pull nvcr.io/nvidia/tensorrt:xx.xx-py3
    2. #第一个xx表示年份，第二个xx表示月份，如22.09，23.01

    成功尝试：直接下载tensorrt的tar包，然后在bin目录下执行trtexec命令，好像是不依赖cuda和cudnn的，但是同样要注意，不同版本的tensorrt编译出来的engine文件，在别的版本tensorrt下可能不能运行。

    使用python接口

    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!")

    使用c++接口

    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<uint32_t>(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. }

    模型推理

    使用python接口

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

    c++接口

    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. }

  • 相关阅读:
    JS定时器实现页面N秒后跳转 实现每隔 1s 自动刷新页面并格式化的显示当前时间
    【Java】监听器
    CiteSpace for Mac 最新保姆级教程
    vivo自研AI大模型即将问世，智能手机行业加速迈向AI时代
    PX4飞行测试
    Stable Diffusion 本地部署教程
    我的测试开发十年之路
    C++学习——C++中const的新花样
    百面深度学习-自然语言处理
    力扣第79题 单词搜索
  • 原文地址：https://blog.csdn.net/zhuikefeng/article/details/133130275