作者: OpenMMLab 原文链接:https://mp.weixin.qq.com/s/ramB73Lr48UJVoHN6ilqHQ 本文仅用于学术分享,如有侵权,请联系后台作删文处理 1. TensorRT 简介 TensorRT 是由 NVIDIA 发布的深度学习框架,用于在其硬件上运行深度学习推理。TensorRT 提供量化感知训练和离线量化功能,用户可以选择 INT8 和 FP16 两种优化模式,将深度学习模型应用到不同任务的生产部署,如视频流、语音识别、推荐、欺诈检测、文本生成和自然语言处理。TensorRT 经过高度优化,可在 NVIDIA GPU 上运行,并且可能是目前在 NVIDIA GPU 运行模型最快的推理引擎。关于 TensorRT 更具体的信息可以访问 TensorRT官网 了解。 TensorRT 官网链接: https://developer./tensorrt 2. 安装 TensorRT Windows 默认在一台有 NVIDIA 显卡的机器上,提前安装好 CUDA 和 CUDNN,登录 NVIDIA 官方网站下载和主机 CUDA 版本适配的 TensorRT 压缩包即可。 以 CUDA 版本是 10.2 为例,选择适配 CUDA 10.2 的 zip 包,下载完成后,有 conda 虚拟环境的用户可以优先切换到虚拟环境中,然后在 powershell 中执行类似如下的命令安装并测试: cd \the\path\of\tensorrt\zip\file Expand-Archive TensorRT-8.2.5.1.Windows10.x86_64.cuda-10.2.cudnn8.2.zip . $env:TENSORRT_DIR = '$pwd\TensorRT-8.2.5.1' $env:path = '$env:TENSORRT_DIR\lib;' + $env:path pip install $env:TENSORRT_DIR\python\tensorrt-8.2.5.1-cp36-none-win_amd64.whl python -c 'import tensorrt;print(tensorrt.__version__)' 上述命令会在安装后检查 TensorRT 版本,如果打印结果是 8.2.5.1,说明安装 Python 包成功了。 zip 包链接: https://developer./compute/machine-learning/tensorrt/secure/8.2.5.1/zip/tensorrt-8.2.5.1.windows10.x86_64.cuda-10.2.cudnn8.2.zip Linux 和在 Windows 环境下安装类似,默认在一台有 NVIDIA 显卡的机器上,提前安装好 CUDA 和 CUDNN,登录 NVIDIA 官方网站下载和主机 CUDA 版本适配的 TensorRT 压缩包即可。 以 CUDA 版本是 10.2 为例,选择适配 CUDA 10.2 的 tar 包,然后执行类似如下的命令安装并测试:
如果发现打印结果是 8.2.5.1,说明安装 Python 包成功了。 tar 包链接: https://developer./compute/machine-learning/tensorrt/secure/8.2.5.1/tars/tensorrt-8.2.5.1.linux.x86_64-gnu.cuda-10.2.cudnn8.2.tar.gz 3. 模型构建 我们使用 TensorRT 生成模型主要有两种方式:
接下来,我们将用 Python 和 C++ 语言分别使用这两种方式构建 TensorRT 模型,并将生成的模型进行推理。 直接构建 利用 TensorRT 的 API 逐层搭建网络,这一过程类似使用一般的训练框架,如使用 Pytorch 或者TensorFlow 搭建网络。需要注意的是对于权重部分,如卷积或者归一化层,需要将权重内容赋值到 TensorRT 的网络中。本文就不详细展示,只搭建一个对输入做池化的简单网络。 使用 Python API 构建 首先是使用 Python API 直接搭建 TensorRT 网络,这种方法主要是利用 此外,需要定义好输入和输出名称,将构建好的网络序列化,保存成本地文件。值得注意的是:如果想要网络接受不同分辨率的输入输出,需要使用 实现代码如下: import tensorrt as trt
verbose = True IN_NAME = 'input' OUT_NAME = 'output' IN_H = 224 IN_W = 224 BATCH_SIZE = 1
EXPLICIT_BATCH = 1 << (int)( trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger() with trt.Builder(TRT_LOGGER) as builder, builder.create_builder_config( ) as config, builder.create_network(EXPLICIT_BATCH) as network: # define network input_tensor = network.add_input( name=IN_NAME, dtype=trt.float32, shape=(BATCH_SIZE, 3, IN_H, IN_W)) pool = network.add_pooling( input=input_tensor, type=trt.PoolingType.MAX, window_size=(2, 2)) pool.stride = (2, 2) pool.get_output(0).name = OUT_NAME network.mark_output(pool.get_output(0))
# serialize the model to engine file profile = builder.create_optimization_profile() profile.set_shape_input('input', *[[BATCH_SIZE, 3, IN_H, IN_W]]*3) builder.max_batch_size = 1 config.max_workspace_size = 1 << 30 engine = builder.build_engine(network, config) with open('model_python_trt.engine', mode='wb') as f: f.write(bytearray(engine.serialize())) print('generating file done!') 使用 C++ API 构建 对于想要直接用 C++ 语言构建网络的小伙伴来说,整个流程和上述 Python 的执行过程非常类似,需要注意的点主要有:
实现代码如下:
IR 转换模型 除了直接通过 TensorRT 的 API 逐层搭建网络并序列化模型,TensorRT 还支持将中间表示的模型(如 ONNX)转换成 TensorRT 模型。 使用 Python API 转换 我们首先使用 Pytorch 实现一个和上文一致的模型,即只对输入做一次池化并输出;然后将 Pytorch 模型转换成 ONNX 模型;最后将 ONNX 模型转换成 TensorRT 模型。 这里主要使用了 TensorRT 的 实现代码如下: import torch import onnx import tensorrt as trt
onnx_model = 'model.onnx'
class NaiveModel(torch.nn.Module): def __init__(self): super().__init__() self.pool = torch.nn.MaxPool2d(2, 2) def forward(self, x): return self.pool(x)
device = torch.device('cuda:0')
# generate ONNX model torch.onnx.export(NaiveModel(), torch.randn(1, 3, 224, 224), onnx_model, input_names=['input'], output_names=['output'], opset_version=11) onnx_model = onnx.load(onnx_model)
# create builder and network logger = trt.Logger(trt.Logger.ERROR) builder = trt.Builder(logger) EXPLICIT_BATCH = 1 << (int)( trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) network = builder.create_network(EXPLICIT_BATCH)
# parse onnx parser = trt.OnnxParser(network, logger)
if not parser.parse(onnx_model.SerializeToString()): error_msgs = '' for error in range(parser.num_errors): error_msgs += f'{parser.get_error(error)}\n' raise RuntimeError(f'Failed to parse onnx, {error_msgs}')
config = builder.create_builder_config() config.max_workspace_size = 1<<20 profile = builder.create_optimization_profile()
profile.set_shape('input', [1,3 ,224 ,224], [1,3,224, 224], [1,3 ,224 ,224]) config.add_optimization_profile(profile) # create engine with torch.cuda.device(device): engine = builder.build_engine(network, config)
with open('model.engine', mode='wb') as f: f.write(bytearray(engine.serialize())) print('generating file done!') IR 转换时,如果有多 Batch、多输入、动态 shape 的需求,都可以通过多次调用 使用 C++ API 转换 介绍了如何用 Python 语言将 ONNX 模型转换成 TensorRT 模型后,再介绍下如何用 C++ 将 ONNX 模型转换成 TensorRT 模型。这里通过 实现代码如下:
4. 模型推理 前面,我们使用了两种构建 TensorRT 模型的方式,分别用 Python 和 C++ 两种语言共生成了四个 TensorRT 模型,这四个模型的功能理论上是完全一致的。 接下来,我们将分别使用 Python 和 C++ 两种语言对生成的 TensorRT 模型进行推理。 使用 Python API 推理 首先是使用 Python API 推理 TensorRT 模型,这里部分代码参考了 MMDeploy。运行下面代码,可以发现输入一个 from typing import Union, Optional, Sequence,Dict,Any
import torch import tensorrt as trt
class TRTWrapper(torch.nn.Module): def __init__(self,engine: Union[str, trt.ICudaEngine], output_names: Optional[Sequence[str]] = None) -> None: super().__init__() self.engine = engine if isinstance(self.engine, str): with trt.Logger() as logger, trt.Runtime(logger) as runtime: with open(self.engine, mode='rb') as f: engine_bytes = f.read() self.engine = runtime.deserialize_cuda_engine(engine_bytes) self.context = self.engine.create_execution_context() names = [_ for _ in self.engine] input_names = list(filter(self.engine.binding_is_input, names)) self._input_names = input_names self._output_names = output_names
if self._output_names is None: output_names = list(set(names) - set(input_names)) self._output_names = output_names
def forward(self, inputs: Dict[str, torch.Tensor]): assert self._input_names is not None assert self._output_names is not None bindings = [None] * (len(self._input_names) + len(self._output_names)) profile_id = 0 for input_name, input_tensor in inputs.items(): # check if input shape is valid profile = self.engine.get_profile_shape(profile_id, input_name) assert input_tensor.dim() == len( profile[0]), 'Input dim is different from engine profile.' for s_min, s_input, s_max in zip(profile[0], input_tensor.shape, profile[2]): assert s_min <= s_input <= s_max, \ 'Input shape should be between ' \ + f'{profile[0]} and {profile[2]}' \ + f' but get {tuple(input_tensor.shape)}.' idx = self.engine.get_binding_index(input_name)
# All input tensors must be gpu variables assert 'cuda' in input_tensor.device.type input_tensor = input_tensor.contiguous() if input_tensor.dtype == torch.long: input_tensor = input_tensor.int() self.context.set_binding_shape(idx, tuple(input_tensor.shape)) bindings[idx] = input_tensor.contiguous().data_ptr()
# create output tensors outputs = {} for output_name in self._output_names: idx = self.engine.get_binding_index(output_name) dtype = torch.float32 shape = tuple(self.context.get_binding_shape(idx))
device = torch.device('cuda') output = torch.empty(size=shape, dtype=dtype, device=device) outputs[output_name] = output bindings[idx] = output.data_ptr() self.context.execute_async_v2(bindings, torch.cuda.current_stream().cuda_stream) return outputs
model = TRTWrapper('model.engine', ['output']) output = model(dict(input = torch.randn(1, 3, 224, 224).cuda())) print(output) MMDeploy 链接: https://github.com/open-mmlab/mmdeploy (欢迎体验,觉得好用欢迎点亮小星星) 使用 C++ API 推理 最后,在很多实际生产环境中,我们都会使用 C++ 语言完成具体的任务,以达到更加高效的代码运行效果,另外 TensoRT 的用户一般也都更看重其在 C++ 下的使用,所以我们也用 C++ 语言实现一遍模型推理,这也可以和用 Python API 推理模型做一个对比。 实现代码如下:
总结 通过本文的学习,我们掌握了两种构建 TensorRT 模型的方式:直接通过 TensorRT 的 API 逐层搭建网络;将中间表示的模型转换成 TensorRT 的模型。不仅如此,我们还分别用 C++ 和 Python 两种语言完成了 TensorRT 模型的构建及推理,相信大家都有所收获!在下一篇文章中,我们将和大家一起学习何添加 TensorRT 自定义算子,敬请期待哦~ FAQ
参考
|
|
来自: mynotebook > 《待分类》