PyTorch转换ONNX，再转换TensorRT，实现模型加速

测试环境介绍

• ubuntu20.04
• 显卡：Tesla 3090
• Driver Version: 460.91.03
• CUDA Version: 11.2

关键代码

pytorch转换ONNX部分：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = BertForSequenceClassification.from_pretrained(out_dir)  
    model.eval().to(device)
    with torch.no_grad():
        input_names = ["input_ids","attention_mask"]
        output_names = ["outputs"]
        input_ids = torch.LongTensor(np.ones((12, 128), dtype=int)).to(device)
        attention_mask = torch.LongTensor(np.ones((12, 128), dtype=int)).to(device)
        torch.onnx.export(
            model = model, 
            args = (input_ids, attention_mask), 
            f = "ai11.onnx",
            opset_version=11,
            input_names=input_names, 
            output_names=output_names)

使用ONNX runtime查看结果是否正确部分：

import onnxruntime
    sim_out_file = 'ai_sim_11.onnx'
    cmd = f"python -m onnxsim ai11.onnx {sim_out_file}"
    print(f'run cmd {cmd}')
    os.system(cmd)

    sess = onnxruntime.InferenceSession("ai11.onnx")
    input_ids = np.ones((12, 128), dtype=int)
    attention_mask = np.ones((12, 128), dtype=int)
    out_onnx = sess.run(None, {
        'input_ids': input_ids,
        'attention_mask': attention_mask
    })
    print(out_onnx)
    pred = torch.tensor(out_onnx)
    pp = torch.nn.functional.softmax(pred, dim=1)
    print(pp)

TensorRT使用部分：

with get_engine(onnx_file_path, engine_file_path) as engine, engine.create_execution_context() as context:
    print('Running inference ...')
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    h_input_1 = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(trt.float32))
    h_input_2 = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(trt.float32))
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(2)), dtype=trt.nptype(trt.float32))
    d_input_1 = cuda.mem_alloc(h_input_1.nbytes)
    d_input_2 = cuda.mem_alloc(h_input_2.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    for batch in prediction_dataloader:
        b_input_ids, b_input_mask = batch

        input_ids = b_input_ids.numpy().astype(trt.nptype(trt.float32)).ravel()
        input_mask = b_input_mask.numpy().astype(trt.nptype(trt.float32)).ravel()
        np.copyto(h_input_1, input_ids)
        np.copyto(h_input_2, input_mask)
        cuda.memcpy_htod_async(d_input_1, h_input_1, stream)
        cuda.memcpy_htod_async(d_input_2, h_input_2, stream)
        context.execute_async_v2(bindings=[int(d_input_1), int(d_input_2), int(d_output)], stream_handle=stream.handle)
        cuda.memcpy_dtoh_async(h_output, d_output, stream)
        stream.synchronize()
        out = torch.tensor(h_output.reshape(12,5))
        out_softmax = torch.nn.functional.softmax(out, dim=1).numpy()
        pred_list = np.argmax(out_softmax, axis=1).flatten()
        print(out)

错误总结及技巧

错误：pytorch转换onnx时候出现各种错误
方法：torch版本改为1.8版本，pip install torch==1.8.0
再次错误：GeForce RTX 3090 with CUDA capability sm_86 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70.
再次方法：pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html  (可以在这个网站查找相应的其他对应版本)

错误：ModuleNotFoundError: No module named 'onnxsim'
方法：pip install onnx-simplifier -i https://pypi.douban.com/simple/

错误：能查到的教程中，onnx导出的模型的输入参数只有一个，nlp中使用bert做推理需要输入input_ids和attention_mask两个参数。
方法：导出模型使args指定元组args = (input_ids, attention_mask)

困难：tensorrt环境部署时候对显卡和驱动版本对应关系要求比较严格，安装部署比较麻烦。
办法：使用nvidia提供的docker环境，部署起来相对来说简单一下，镜像地址为：https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt，如果需要查询历史版本和对应的环境映射，点击其中的链接“TensorRT Container Release Notes”

错误：docker: Error response from daemon: OCI runtime create failed: unable to retrieve OCI runtime error (open /run/containerd/io.containerd.runtime.v2.task/moby/61fd8a8c0ccfe6912104b0198e636ba8755ed7066cbd0e6385dcbc10d3301794/log.json: no such file or directory): fork/exec /usr/bin/nvidia-container-runtime: no such file or directory: unknown.
办法：sudo apt-get install nvidia-container-runtime

错误：容器里面安装依赖包连不上网（启动容器时指定使用本机网卡）
方法：启动容器时候加上参数network，docker run -it --network=host imageid /bin/bash

错误：No importer registered for op: NonZero
分析：这个问题是TensorRT还没有支持ONNX中的算子NonZero（https://github.com/onnx/onnx-tensorrt/issues/401），两个方法避免：1、自己写算子加入到TensorRT中；2、从模型中避免这个算子的应用。我查看本机pytorch中的onnx模型，发现这个算子在opset9和13中有（https://github.com/onnx/onnx/blob/master/docs/Operators.md#NonZero），在其他opset中没有，我猜测在其他opset中是使用了其他替换方式，所以pytoch转换onnx时候指定opset版本可以避免这个算子的出现。
方法：pytorch转换onnx时候指定opset_version=11

困难：由于使用C++进行数据的前后处理的话，有一些功能使用C++复现暂时成本过高，所以采用python来做，并且docker容器中自带的case处理的单输入比较多，我们使用多输入。
方法：上线推理功能时候，仍旧使用pytorch中的TensorDataset和DataLoader封装数据流，参考docker容器中自带的case调用tensorrt模型进行推理。

错误：使用tensorrt加速后计算的结果和未使用加速的结果不相同。
分析：最终发现不使用加速时候输入数据类型是int32，而tensorrt方案使用的输入是float32,类型转换发生了数据变化。
方案：tensorrt解决方案输入数据类型由trt.float32转化为trt.int32！


技巧：
1、在容器和本地之间互拷内容
docker cp 83bd16c2b7bf:/workspace/tensorrt/samples/python/yolov3_onnx/onnx_to_tensorrt.py ./