【TensorRT】PyTorch模型转换为ONNX及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账号。

• 新建文件夹，将压缩文件拷贝进来解压：

tar xzvf TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

• 解压得到TensorRT-8.4.3.1的文件夹，将里面的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命令安装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 range(len(input_names)):
        dims = engine.get_profile_shape(0, binding_idx)[2] # 拿到max的dim
        count = 1
        for dim in dims:
            count *= dim
        input_mem = count * 4
        # 错误: explicit_context_dependent failed: invalid device context - no currently active context?
        # import pycuda.autoinit 即可
        buffers.append(cuda.mem_alloc(input_mem))  
    
    # # 分配输出的内存
    for binding_idx in range(len(input_names), len(input_names) + len(output_names)):
        binding_name = engine.get_binding_name(binding_idx)
        dims = output_shapes[binding_name]
        count = 1
        for dim in dims:
            count *= dim
        input_mem = count * 4
        buffers.append(cuda.mem_alloc(input_mem))  
    
    pred_binding_idx = engine.get_binding_index(output_names[output_idx])  # 预测轨迹的binding index
    return context, buffers, pred_binding_idx

vector_data，vector_mask，traj_polyline_mask，map_polyline_mask，cent_polyline_mask为输入，供参考。

TensorRT 推理：

def 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):
    """inference
    Args:
        vector_data (tensor)
        vector_mask (tensor)
        traj_polyline_mask (tensor)
        map_polyline_mask (tensor)
        device (torch.device)
        num (int): loop num
    Returns:
        tensor: pred traj
    """
    # 数据预处理
    all_polyline_idx, vector_idx, invalid_idx, map_polyline_idx, traj_polyline_idx, cent_polyline_idx, topk_points_idx = \
        data_preprocess(vector_mask, traj_polyline_mask, map_polyline_mask, cent_polyline_mask, device)

    trt_inputs = (vector_data.to(device).numpy(), all_polyline_idx.to(torch.int32).numpy(), vector_idx.to(torch.int32).numpy(), 
                    invalid_idx.to(torch.int32).numpy(), map_polyline_idx.to(torch.int32).numpy(), traj_polyline_idx.to(torch.int32).numpy(), 
                    cent_polyline_idx.to(torch.int32).numpy(), topk_points_idx.to(torch.int32).numpy())  # .astype(np.float64)

    # # 分配内存空间，并进行数据cpu到gpu的拷贝
    # 动态尺寸，每次都要set一下模型输入的shape, 由于输出shape是固定的，只需要set输入的shape
    for i in range(len(trt_inputs)):
        context.set_binding_shape(i, trt_inputs[i].shape)
    
    # # 将输入数据从cpu拷贝到gpu
    for binding_idx in range(len(trt_inputs)):
        cuda.memcpy_htod(buffers[binding_idx], trt_inputs[binding_idx])  
    
    # # 进行推理，并将结果从gpu拷贝到cpu
    output = np.empty(output_shapes[output_names[output_idx]], dtype=np.float32)
    context.execute_v2(buffers)  # 可异步和同步
    cuda.memcpy_dtoh(output, buffers[pred_binding_idx])
    return output

5. 转换TensorRT需要注意的一些语法规则

TensorRT 有一定的算子限制，TensorRT支持的ONNX算子点击查看。

在改写 PyTorch 代码以转换 ONNX 时有一些特殊的语法限制，转换 TorchScript 的语法规则同样适用于 TensorRT，但 TensorRT 的语法限制更为严格。

5.1 tensor索引不支持bool类型作为索引参数

Python中使用bool数据进行index索引：

vector_data = torch.randn(3,4,5)
map_polyline_mask = tensor([False, True, True])
print(vector_data)
print(vector_data[map_polyline_mask])

输出：

tensor([[[ 0.5749, -0.0499,  0.3597,  1.6321, -0.8703],
         [ 0.8106,  0.0916, -0.2107, -0.2349, -0.4759],
         [-0.4364,  0.6143, -2.3013,  0.2073, -0.5179],
         [ 1.3962,  0.9612, -1.1123, -0.3694, -1.9012]],

        [[-0.1141,  0.3616, -0.0526, -0.7861, -0.3961],
         [ 0.1197, -0.3391,  0.3159, -0.6221,  0.4027],
         [-0.4965, -0.5303,  0.0416,  1.9262,  0.2784],
         [ 0.0797, -0.4318,  0.8028,  0.2495,  0.8283]],

        [[-0.0132,  1.1342, -0.0522,  0.5333,  0.0553],
         [ 0.1323,  0.0501, -1.7931,  0.4482,  0.0596],
         [ 0.2548,  0.6196, -1.1362, -0.4189,  1.2865],
         [ 0.2258, -1.9663, -1.0206, -0.2433,  0.6593]]])

tensor([[[-0.1141,  0.3616, -0.0526, -0.7861, -0.3961],
         [ 0.1197, -0.3391,  0.3159, -0.6221,  0.4027],
         [-0.4965, -0.5303,  0.0416,  1.9262,  0.2784],
         [ 0.0797, -0.4318,  0.8028,  0.2495,  0.8283]],

        [[-0.0132,  1.1342, -0.0522,  0.5333,  0.0553],
         [ 0.1323,  0.0501, -1.7931,  0.4482,  0.0596],
         [ 0.2548,  0.6196, -1.1362, -0.4189,  1.2865],
         [ 0.2258, -1.9663, -1.0206, -0.2433,  0.6593]]])

Python中可以正常使用bool数据进行index索引，TensorRT中无法使用bool数据进行index索引。

解决方法： 将bool索引改为int索引。

修改示例：

map_polyline_idx = torch.vstack(torch.where(map_polyline_mask==1))[0]
print(map_polyline_idx)
print(vector_data[map_polyline_idx])

输出：

tensor([1, 2])

tensor([[[-0.1141,  0.3616, -0.0526, -0.7861, -0.3961],
         [ 0.1197, -0.3391,  0.3159, -0.6221,  0.4027],
         [-0.4965, -0.5303,  0.0416,  1.9262,  0.2784],
         [ 0.0797, -0.4318,  0.8028,  0.2495,  0.8283]],

        [[-0.0132,  1.1342, -0.0522,  0.5333,  0.0553],
         [ 0.1323,  0.0501, -1.7931,  0.4482,  0.0596],
         [ 0.2548,  0.6196, -1.1362, -0.4189,  1.2865],
         [ 0.2258, -1.9663, -1.0206, -0.2433,  0.6593]]])

修改之后，转TensorRT通过。

5.2 squeeze()会导致ONNX模型出现 if 节点

原因是 squeeze() 仅压缩shape为1的维度，如果压缩的维度shape不是1则不会压缩。

解决方法： 使用 .view() 来替换squeeze() 。

unsqueeze() 则不受影响，因为扩展维度总是能扩展的。

5.3 argmax需要至少2维的tensor作为输入

例：

scores = torch.tensor([-10, -1, 2, 5])
torch.argmax(scores)

上述代码转换为TensorRT会报需要至少2维的tensor，替换方案：

torch.argmax(scores.unsqueeze(0), 1, False)

5.4 expand操作转换TensorRT问题

expand的shape如果是从输入传入的动态数值，可以正常转换TensorRT模型，如果onnx能通过推理得到expand的shape那onnx模型中该shape会被记录为常数，转TensorRT就可能出问题。

5.5 repeat操作转换TensorRT问题

TensorRT支持repeat操作，但repeat传入的参数必须是固定的，比如：valid_mask_new.repeat(batch_size, 1, 1, 1) 中，如果batch size的值是固定的，转换TensorRT没有问题，但如果batch size的值是动态的，转换TensorRT则会失败。

5.6 dimensions not compatible for scatterND

scatterND问题：

valid_mask = torch.zeros(5)
vector_idx = torch.tensor([0, 1, 4])
valid_mask[vector_idx] = float(1)
print(valid_mask)

输出：

tensor([1., 1., 0., 0., 1.])

Python中索引赋值正常，转TensorRT报错：dimensions not compatible for scatterND 。

官方给出的scatterND案例：

data    = [1, 2, 3, 4, 5, 6, 7, 8]
indices = [[4], [3], [1], [7]]
updates = [9, 10, 11, 12]
output  = [1, 11, 3, 10, 9, 6, 7, 12]

可见 data 的 shape 为[8]，indices 的 shape为[4, 1]，indices 比 data 多了一维。

解决方法： 使用 unsqueeze 增加 index 的维度。

加入unsqueeze：

valid_mask[vector_idx.unsqueeze(-1)] = float(1)
print(valid_mask)

输出：

tensor([1., 1., 0., 0., 1.])

加入unsqueeze 后，转TensorRT通过。

5.7 数据类型为int32的多维input传入TensorRT与原数据不等

做了几次测试，输入TensorRT的数据为2维的int32型的tensor，将输入直接作为output进行输出，结果均与输入不等，原因未知。

Python 中正确的2维 index索引：

valid_mask_new[vector_idx[:,0].unsqueeze(1), vector_idx[:,1].unsqueeze(1)] = float(1)

另外2维的 index 索引在 Python 中只能使用上述方式，转换 TensorRT 后该索引方式结果与 Python 不一致，所以尽量避免2维的 index 索引。

解决方法：

将多维的index输入flatten到1维，修改对应位置的数值，同样模型内部gather的原数据也需要flatten，使用index索引完之后再将原数据view到原来的维度即可。

index数据flatten如：

vector_idx = (vector_idx[0, :, 0] * 20 + vector_idx[0, :, 1]).unsqueeze(0)  # index维度的长度不能超过1

5.8 定位TensorRT的错误fusion，并拆分错误fusion

转换 TensorRT 会将 onnx 模型内部的很多 node 进行融合以加速推理，这一部分 TensorRT 并没有开源，所以如何进行融合的不得而知。只能通过将各个 node 的 output 加入模型的 output 进行输出，通过二分法来拆分错误的 fusion。

import onnx

onnx_model = onnx.load("trt/dense_tnt_sim.onnx")
for value_info in onnx_model.graph.value_info:
	id = eval(value_info.name)
	if 1688 <= id <= 1751 and value_info.type.tensor_type.elem_type == 1:  # 只加入float类型的value作为输出
	    out1 = onnx.helper.make_tensor_value_info(str(id), onnx.TensorProto.DataType.FLOAT, (764,764))
	    # out1 = onnx.ValueInfoProto(name=str(id))
	    onnx_model.graph.output.append(out1)
	    idx += 1

5.9 TensorRT不支持torch.topk()动态k值

torch.topk()这个算子TensorRT是支持的，但不管k值是模型内部的动态数值，还是输入将k值作为输入传入TensorRT，前者会将k值固化为常数，后者会在onnx模型中记录为动态数值，但转TensorRT会报无法接收动态k值的错误。

尝试的方法：将topk改为sort排序，然后使用slice操作取出前k（有k条数据就取前k，没有则取所有，间接实现了topk操作），但转为onnx后发现onnx将sort和slice操作转换成了topk，故该方法无效。

注意： 由于k值只能是常数，所以k值的设定和输入数据两方面都需要注意，保证模型内计算topk的原数据的条数大于等于k，否则TensorRT计算会导致core dump。