文章目录
- 1. PyTorch模型转TensorRT模型流程
- 2. PyTorch模型转ONNX模型
- 3. ONNX模型转TensorRT模型
- 4. TensorRT在Python中推理
- 5. 转换TensorRT需要注意的一些语法规则
-
- 5.1 tensor索引不支持bool类型作为索引参数
- 5.2 squeeze()会导致ONNX模型出现 if 节点
- 5.3 argmax需要至少2维的tensor作为输入
- 5.4 expand操作转换TensorRT问题
- 5.5 repeat操作转换TensorRT问题
- 5.6 dimensions not compatible for scatterND
- 5.7 数据类型为int32的多维input传入TensorRT与原数据不等
- 5.8 定位TensorRT的错误fusion,并拆分错误fusion
- 5.9 torch.topk()相关问题
- 5.10 TensorRT对int64支持较差(trt版本:8.4.3.1)
- 5.11 torch.scatter_reduce()相关问题
- 5.12 repeat替换为低时延算子
- 5.13 pow()替换为低时延算子
- 5.14 PWN算子优化
- 6. TensorRT模型耗时分析
任务简介:
TensorRT 模型的推理速度比 libtorch 模型更快,所以 PyTorch 模型转换为 TensorRT 模型部署几乎是最好的选择。通常TensorRT 模型首先需要转换为 ONNX 模型,再由 ONNX 模型转换为TensorRT 模型。本文对转换方法及一些注意点做一个记录。
1. PyTorch模型转TensorRT模型流程
2. PyTorch模型转ONNX模型
pytorch模型转onnx模型示例:
torch.onnx.export(model, # model being run
onnx_inputs, # model input (or a tuple for multiple inputs)
"trt/dense_tnt.onnx", # where to save the model (can be a file or file-like object)
export_params=True, # store the trained parameter weights inside the model file
verbose=True, # if True, all parameters will be exported
opset_version=11, # the ONNX version to export the model to
do_constant_folding=False, # whether to execute constant folding for optimization
input_names=['vector_data', 'all_polyline_idx', 'vector_idx', 'invalid_idx', 'map_polyline_idx', 'traj_polyline_idx', 'cent_polyline_idx', 'topk_points_idx'], # the model's input names
output_names=['trajectories'], # the model's output names
dynamic_axes={
'all_polyline_idx': {
1: 'all_polyline_num'},
'vector_idx': {
1: 'vector_num'},
'invalid_idx': {
1: 'invalid_num'},
'map_polyline_idx': {
1: 'map_polyline_num'},
'traj_polyline_idx': {
1: 'traj_polyline_num'},
'cent_polyline_idx': {
1: 'cent_polyline_num'},
'topk_points_idx': {
1: 'topk_points_num'},})
verbose=True
可以将转换后的模型代码及参数输出,并对应了相应的源代码。
dynamic_axes
可以设置动态输入,如 'all_polyline_idx': {1: 'all_polyline_num'}
表示 all_polyline_idx 的第1维的shape为动态,且命名为 all_polyline_num。
ONNX 模型简化:
onnxsim input_onnx_model_name output_onnx_model_name
简化后输出:
3. ONNX模型转TensorRT模型
3.1 TensorRT安装
- 点击到官网下载,选择自己需要的版本,需要nvidia账号。
- Install TensoRT
tar -xzvf TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
sudo cp TensorRT-8.4.3.1/include/* /usr/include/x86_64-linux-gnu/
sudo cp -r TensorRT-8.4.3.1/lib/* /usr/lib/x86_64-linux-gnu/
sudo cp TensorRT-8.4.3.1/bin/* /usr/bin/
- (可选)将里面的lib绝对路径添加到环境变量中:
export LD_LIBRARY_PATH=${TENSORRT_PATH}/TensorRT-8.4.3.1/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=${TENSORRT_PATH}/TensorRT-8.4.3.1/lib:$LIBRARY_PATH
- 使用pip命令安装python版本的TensorRT:
cd TensorRT-8.4.3.1/python/
pip install tensorrt-8.4.3.1-cp38-none-linux_x86_64.whl
3.2 将ONNX模型转换为TensorRT模型
转换命令:
${TENSORRT_PATH}/TensorRT-8.4.2.4/bin/trtexec
--onnx=dense_tnt_sim.onnx
--minShapes=all_polyline_idx:1x7,vector_idx:1x6,invalid_idx:1x1,map_polyline_idx:1x6,traj_polyline_idx:1x1,cent_polyline_idx:1x3,topk_points_idx:1x50
--optShapes=all_polyline_idx:1x1200,vector_idx:1x24000,invalid_idx:1x24000,map_polyline_idx:1x1200,traj_polyline_idx:1x1200,cent_polyline_idx:1x1200,topk_points_idx:1x24000
--maxShapes=all_polyline_idx:1x1200,vector_idx:1x24000,invalid_idx:1x24000,map_polyline_idx:1x1200,traj_polyline_idx:1x1200,cent_polyline_idx:1x1200,topk_points_idx:1x24000
--saveEngine=dense_tnt_fp32.engine
--device=0
--workspace=48000
--noTF32
--verbose
--minShapes
为inputs的最小维度;
--optShapes
为输入常用的inputs维度,我这边输入的是最大维度;
--maxShapes
为inputs的最大维度;
--device
设置转换模型使用的gpu;
--noTF32
不使用tf32数据类型,使用fp32;
--verbose
输出详细信息。
4. TensorRT在Python中推理
整体流程:
if __name__ == '__main__':
device = torch.device("cpu") # TensorRT模型不管device是cpu还是cuda都会调用gpu,为了方便转换用cpu可以解决device冲突问题
trt_path = "/home/chenxin/peanut/DenseTNT/trt/densetnt_vehicle_trt_model_tf32.engine"
input_names = ["vector_data", "all_polyline_idx", "vector_idx", "invalid_idx", "map_polyline_idx", "traj_polyline_idx",
"cent_polyline_idx", "topk_points_idx"]
output_names = ["1195", "1475", "1785"] # 1836 is pred traj
output_shapes = {
"1195": (1, 1200, 128), "1475": (1, 1200, 128), "1785": (1, 30, 2)}
output_idx = 2
# 1. 初始化模型, 分配输入输出的内存
context, buffers, pred_binding_idx = RTEngineInit(trt_path, input_names, output_names, output_shapes, output_idx)
# 2. 模型推理
pred = trt_inference(context, buffers, output_names, output_shapes, output_idx, pred_binding_idx, vector_data, vector_mask, traj_polyline_mask, map_polyline_mask, cent_polyline_mask, device)
pred = torch.tensor(pred)
print(pred)
加载TensorRT engine文件:
def load_engine(engine_path):
TRT_LOGGER = trt.Logger(trt.Logger.INFO) # INFO
# TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
with open(engine_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
初始化 TensorRT 模型:
def RTEngineInit(trt_path, input_names, output_names, output_shapes, output_idx):
# 建立模型,构建上下文管理器
engine = load_engine(trt_path)
# 创建context用来执行推断
context = engine.create_execution_context()
context.active_optimization_profile = 0
buffers = []
# # 分配输入的内存
for binding_idx in