1. Tensorrt加速原理
1、Precision Calibration
精度校准——训练时由于梯度等对于计算精度要求较高，但是inference阶段可以利用精度较低的数据类型加速运算，降低模型的大小，例如FP16，int8，从而加速模型推理速度。
2、Layer & Tensor fusion
层和张量融合——TensorRT中将多个层的操作合并为同一个层，这样就可以一定程度的减少kernel launches和内存读写。比如把主流神经网络的conv、BN、Relu三个层融合为了一个层；把维度相同的张运算组合成另一个大的张量运算。每一层的运算操作都是由GPU完成的——GPU通过启动不同的CUDA（Compute unified device architecture）核心来完成计算的，CUDA核心计算张量的速度是很快的，但是往往大量的时间是浪费在CUDA核心的启动和对每一层输入/输出张量的读写操作上面，这造成了内存带宽的瓶颈和GPU资源的浪费。
3、Kernel Auto-Tuning
计算核心自动调整——TensorRT可以针对不同的算法，不同的网络模型，不同的GPU平台，进行 CUDA核的调整，以保证当前模型在特定平台上以最优性能计算。
4、 Dynamic Tensor Memory
动态张量显存——每个tensor的使用期间，TensorRT会为其指定显存，避免显存重复申请，减少内存占用和提高重复使用效率。
5、Multi-Stream Execution
并行处理多流输入——这个是GPU底层优化，理解不了。
6、 Time Fusion
时间融合——使用动态生成的算子优化循环神经网络。
以上对为何能加速进行了简单的介绍，详细的原理很难有比较深刻的理解。总体就是量化——降低数据精度、cuda kernel 智能化计算、动态显存管理以及模型结构和张量融合突破GPU带宽瓶颈。
2. 安装tensorrt
官方文档： Quick Start Guide :: NVIDIA Deep Learning TensorRT Documentation
本文安装的Tensorrt为8.6.0 版本，建议安装最新的稳定版本。
环境 ： cuda11.4， python3.8， 1080Ti显卡，Linux ubuntu。
3. pytorch -->onnx
onnx分静态输入和动态输入。动态输入在转trt的时候相对复杂一点。由于nnunet patch size固定。所以可转静态输入模型。可参考https://zhuanlan.zhihu.com/p/548006090#pytorch%E6%A8%A1%E5%9E%8B%E8%BD%AC%E4%B8%BA%E4%BB%85%E6%94%AF%E6%8C%81%E9%9D%99%E6%80%81%E8%BE%93%E5%85%A5%E7%9A%84onnx%E6%A8%A1%E5%9E%8B
动态输入转化代码如下：
net = trainer.network
checkpoint = torch.load(os.path.join(Model_path, folds[0], checkpoint_name + ".model"))
net.load_state_dict(checkpoint['state_dict'])
net.eval()
dummy_input = torch.randn(2, 1, 128, 128, 128).to("cuda")
torch.onnx.export(
net,
dummy_input,
'dynamic_nnunet_3.onnx',
input_names=['input'],
output_names=['output'],
dynamic_axes={ # 动态输入，表示0，1，2，3，4输入都可变，
# 此时dummy_input数值可以随意输入，batchsize为1可能会报错
'input': [0, 1, 2, 3, 4],
'output': [0, 1, 2]}
)静态输入转化代码：
dummy_input = torch.randn(1, 1,96, 160, 160).to("cuda") # 静态输入这里必须和输入图像尺寸一致
torch.onnx.export(
net,
dummy_input,
'static_nnunet_3.onnx',
input_names=['input'],
output_names=['output']
)转出onnx模型验证：
Model = onnx.load(Onnx_file)
onnx.checker.check_model(Model) # 验证Onnx模型是否准确
netron.start('/to_onnx/file/dynamic_nnunet_3.onnx') # 可视化转出模型结构部分模型在转化的时候会由于onnx没有实现相关函数报错，只能修改对应源码或者等新版本发布。
可以用onnxruntime来检测转出onnx模型的正确性，以下为onnxruntime推理的简单示例。nnUnet 修改主要修改nnunet/inference/predict.py > predict_cases() > predict_preprocessed_data_return_seg_and_softmax() > predict_3D() >_internal_predict_3D_3Dconv_tiled 这个函数
import os
import onnxruntime as rt
import numpy as np
import torch
from nnunet.network_architecture.neural_network import SegmentationNetwork
from pre.utils import get_img_array, save_processed
file_path = '/home/youjiangchuan/Documents/coronary_220_0000.nii.gz'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ia, o,d,s = get_img_array(file_path)
data = np.expand_dims(np.expand_dims(ia, axis=0), axis=0)
# model
sess = rt.InferenceSession('/to_onnx/file/nnunet_pas.onnx',
providers=['CUDAExecutionProvider'])
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
steps = SegmentationNetwork._compute_steps_for_sliding_window(patch_size=(96,160,160),image_size=data[0, 0,:,:,:].shape, step_size=0.5)
# [[0, 40, 81, 121, 162, 202], [0, 70, 141, 211, 282, 352], [0, 70, 141, 211, 282, 352]]
patch_size = (96,160,160)
class_num = 2
aggregated_results = np.zeros([1,class_num] + list(data[0,0,:,:,:].shape)).astype(np.float32)
print(aggregated_results.shape)
for x in steps[0]:
lb_x = x
ub_x = x + patch_size[0]
for y in steps[1]:
lb_y = y
ub_y = y + patch_size[1]
for z in steps[2]:
lb_z = z
ub_z = z + patch_size[2]
predicted_patch = sess.run([label_name], {input_name: data[:,:, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z].astype(np.float32)})[0]
aggregated_results[0, 0, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z] += predicted_patch[0, 0, :, :, :]
aggregated_results[0, 1, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z] += predicted_patch[0, 1, :, :, :]
# predict
pred = np.squeeze(aggregated_results)
save_processed(np.argmax(pred, axis=0)[:96,:160,:160], o, d, s, '../file/pre_patch_220.nii.gz')可以用np.testing.assert_allclose 来对比torch和onnx的预测结果误差。
4. onnx --> tensorrt
将pytorch训练好的模型转换成onnx模型再转化为tensorrt下的.trt文件，有两种方法
-
采用TensorRT自带的trtexec把onnx转化为.trt
-
调用tensorrt的API把onnx转化为.trt
4.1 采用trtexec转化
采用trt自带的trtexec转化， 命令为 ./trtexec --onnx= "onnx path " --saveEngine='nnUnet.egnine' --explicitBatch --inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16
./trtexec --h可以查看全部的参数。根据错误信息调整命令，上述命令为静态输入转化命令。成功转化显示如下：
转化日志中会显示很多信息，GPU参与运算，会给出不同batch size的平均推理时间，吞吐量之类信息。可以仔细看一下。
还有一种更直接的办法就是使用tensorrt的API直接把pytorch下的bert模型构建成一个trt的engine执行推理，这个需要对tensorrt的API非常熟悉，以及bert模型的权重和结构非常熟悉，对代码功底要求也比较高。尝试了一下动态的没成功，转出engine为空，静态输入没有试。具体可以参考官方sample/python中的例子，以及上述教程https://zhuanlan.zhihu.com/p/548006090#pytorch%E6%A8%A1%E5%9E%8B%E8%BD%AC%E4%B8%BA%E4%BB%85%E6%94%AF%E6%8C%81%E9%9D%99%E6%80%81%E8%BE%93%E5%85%A5%E7%9A%84onnx%E6%A8%A1%E5%9E%8B
4.2 采用API转化
主要的API是：
trt.Builder(TRT_LOGGER) as builder创建一个builder
builder.create_network(explicit_batch) as network创建一个空的网络
trt.OnnxParser(network, TRT_LOGGER) as parser创建一个onnx解析器
trt.Runtime(TRT_LOGGER) as runtime创建一个trt的运行环境
config = builder.create_builder_config()创建builder_config
设置一些属性后，就可以直接解析onnx然后转化为trt engine
plan = builder.build_serialized_network(network, config)
engine = runtime.deserialize_cuda_engine(plan)5. 用TensorRT进行推理加速
5.1 重要概念说明
-
Logger：日志记录器，可用于记录模型编译的过程
-
Builder：可用于创建 Network，对模型进行序列化生成engine
-
Network：由 Builder 的创建，最初只是一个空的容器
4. Parser：用于解析 Onnx 等模型
5. context：上接 engine，下接 inference，因此解释为上下文
5.2 流程说明
-
基于 Tensor RT 进行模型的推理，从宏观上可分为三个阶段，即输入数据的前处理（preprocess）、模型推理（inference）和推理结果的后处理（postprocess）；
-
在模型推理中，一般会将 Onnx 等格式的模型通过编译（build）生成推理引擎（engine），engine 可以通过序列化（serialize）的方式进行永久存储，永久存储的 engine 则可以通过反序列化（deserialize）的方式进行加载；
-
在正式推理（inference）前，需要手动申请 Cuda 的内存（memory），以存储输入和输出数据流，数据将以流（stream）的形式进行传递；推理前，输入数据（input_data）需要从主机（host）转移到 Cuda（device）中，推理结束后的推理结果则需要从 Cuda（device）转移到主机（host）的内存当中；
-
编译模型时，必须需要定义 logger，并使用定义的 parser 进行解析模型，config用于配置模型，如通过 profile 设定模型的动态输入尺寸，通过 builder 可以创建序列化的 engine；
-
序列化的 engine 需要经过反序列化，才能用于创建 context，最后进行推理
5.3 重要函数说明
-
对于每一个输入张量与输出张量，都需要分配两块资源，分别是主机内存（Host）中的资源以及显存（Device）中的资源。
-
在主机内存（Host）中分配空间，使用 pycuda.driver.cuda.pagelocked_empty(shape, dtype)。shape 一般通过 trt.volume(engine.get_binding_shape(id))实现，可以理解为元素数量（而不是内存大小）。dtype就是数据类型，可以通过 np.float32 或 trt.float32 的形式。
-
显存（Device）中分配空间，使用 pycuda.driver.cuda.mem_alloc(buffer.nbytes)， buffer 可以是ndarray，也可以是前面的 pagelocked_empty() 结果。
-
数据从Host拷贝到Device，使用 pycuda.driver.cuda.memcpy_htod(dest, src)，dest是 mem_alloc 的结果，src 是 numpy/pagelocked_empty。
-
数据从Device拷贝到Host，使用 pycuda.driver.cuda.memcpy_dtoh(dest, src)，dest是numpy/pagelocked_empty，src是mem_alloc。
-
binding可以理解为端口，表示 input tensor与 output tensor，可通过 id 或 name 获取对应的 binding。在模型推理过程中，需要以 bindings 作为输入，其具体数值为内存地址，即 int(buffer)。
-
bindings是一个数组，包含所有的input/output buffer（即device）的地址，获取方式就是直接通过 int(buffer)，其中 buffer 就是 mem_alloc 的结果。
5.4 nnUnet用tensorrt加速例子
更多例子可以看官方给的sample但是官方sample用的是cuda，cudart，本文用的是pycuda。暂不清楚两者之间差异。
从engine中获取inputs, outputs, bindings, stream的格式以及分配缓存.
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None):
"""
分配buffer memory， 存放input 和 output
"""
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)] # ['input', '143', '152', '161', '170', 'output']
# print(tensor_names)
# bind 绑定的输入输出
for binding in tensor_names:
# get_tensor_profile_shape returns (min_shape, optimal_shape, max_shape)
# Pick out the max shape to allocate enough memory for the binding.
shape = engine.get_tensor_shape(binding) if profile_idx is None else \
engine.get_tensor_profile_shape(binding, profile_idx)[-1]
# print(engine.get_binding_shape(binding))
# print(engine.get_tensor_shape(binding))
# (1, 1, 96, 160, 160)
# (1, 2, 6, 10, 10)
# (1, 2, 12, 20, 20)
# (1, 2, 24, 40, 40)
# (1, 2, 48, 80, 80)
# (1, 2, 96, 160, 160)
# shape_valid = np.all([s >= 0 for s in shape]) 检查shape是否有效
size = trt.volume(shape) # binding 的 size
# print(size)
if engine.has_implicit_batch_dimension:
size *= engine.max_batch_size
dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(binding)))
# print(dtype)
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, streamTensorrt 推理
def do_inference_v2(context, bindings, inputs, outputs, stream):
"""
input data (CPU) --> GPU --> run inference (GPU) --> CPU
"""
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async_v2(bindings = bindings, stream_handle= stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]获取engine，建立上下文。
可能会出现以下错误：
-
getPluginCreator could not find plugin 【找不到的OP名称】 version 1
原因 ：对应op放到了插件库，需加载插件库
解决方法：
# 在反序列化前加上，trt_runtime.deserialize_cuda_engine(engine)
trt.init_libnvinfer_plugins(None, "")if __name__ == '__main__':
file_path = ''
ia, o, d, s = get_img_array(file_path)
data = np.expand_dims(np.expand_dims(ia, axis=0), axis=0)
# model
steps = compute_steps_for_sliding_window(patch_size=(96, 160, 160), image_size=data[0, 0, :, :, :].shape,
step_size=0.5)
patch_size = (96, 160, 160)
class_num = 2
aggregated_results = np.zeros([1, class_num] + list(data[0, 0, :, :, :].shape)).astype(np.float16)
# 这句一定要加，不加会报错
trt.init_libnvinfer_plugins(None, "")
f = open("/home/youjiangchuan/PycharmProjects/Tensorrt/nnUnet_s_egnine.trt", "rb")
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(f.read())
# 推理环境上下文
context = engine.create_execution_context()
inputs, outputs, bindings, stream = allocate_buffers(engine)
for x in steps[0]:
lb_x = x
ub_x = x + patch_size[0]
for y in steps[1]:
lb_y = y
ub_y = y + patch_size[1]
for z in steps[2]:
lb_z = z
ub_z = z + patch_size[2]
inputs[0].host = np.ascontiguousarray(data[:, :, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z], dtype=np.float16)
trt_oututs = do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
aggregated_results[0, 0, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z] += trt_oututs[-1].reshape(
(1, 2, 96, 160, 160))[0, 0, :, :, :]
aggregated_results[0, 1, lb_x:ub_x, lb_y:ub_y, lb_z:ub_z] += trt_oututs[-1].reshape(
(1, 2, 96, 160, 160))[0, 1, :, :, :]
# predict
pred = np.squeeze(aggregated_results)
save_processed(np.argmax(pred, axis=0), o, d, s, 'trt_patch_220.nii.gz')6. 结果
tensorrt预测一例在1080Ti需要113s， pytorch 需要369s，速度提高了3倍左右。精度误差在0.002%左右。
除了fp16还可以用 int8，int8比较麻烦的应该是需要一个精度校准器Calibrator.
动态输入在每个环节都略有不同。
扫一扫
9551
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
