安装
按照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过来…