TensorRT尝鲜

安装

按照quick-start-guid安装方法来装就行, 这里推荐使用docker镜像安装.

下载链接

命令:

docker pull nvcr.io/nvidia/tensorrt:21.06-py3

不过cuda的版本要和驱动匹配起来, 直接从docker容器名称也看不错这个镜像里装得到底是那个版本的cuda, workaround的方法是看下layers页面中的cuda相关字段.

安装完成后启动该镜像:

xhost +local:${USER}
docker run --gpus all -it --rm     \
    -e DISPLAY=$DISPLAY     \
    -u $(id -u) \
    -v /tmp/.X11-unix:/tmp/.X11-unix     \
    -v ${PWD}:${PWD} \
    -w ${PWD} \
    --name tensorrt \
    nvcr.io/nvidia/tensorrt:21.06-py3

注意docker run后面要加上 —gpus all 选项.
进入容器后可以转换模型了

wget https://s3.amazonaws.com/download.onnx/models/opset_8/resnet50.tar.gz
tar xzf resnet50.tar.gz
trtexec --onnx=resnet50/model.onnx --saveEngine=resnet_engine.trt  --explicitBatch

• 安装pycuda
  由于ONNXClassifierWrapper类中使用了pycuda, 所以在运行时还要额外安装.不过我使用的conda隔离环境,build时提示找不到cuda.h, 网上的一种解决方法是:

pip3 install --global-option=build_ext --global-option="-I/usr/local/cuda-10.0/targets/aarch64-linux/include/" --global-option="-L/usr/local/cuda-10.0/targets/aarch64-linux/lib/" pycuda

然而并没有用.
最后使用

conda install -c conda-forge pycuda

安装成功…

测试

import tensorrt as trt
def demo1():
    import numpy as np
    PRECISION = np.float32
    from onnx_helper import ONNXClassifierWrapper

    BATCH_SIZE = 1
    N_CLASSES = 1000 # Our ResNet-50 is trained on a 1000 class ImageNet task
    trt_model = ONNXClassifierWrapper("resnet_engine.trt", [BATCH_SIZE, N_CLASSES], target_dtype = PRECISION)
    dummy_input_batch = np.zeros((BATCH_SIZE, 224, 224, 3))
    predictions = trt_model.predict(dummy_input_batch)

注意上面的onnx_helper要从github上的tensorrt仓库中copy过来…