创建conda环境
conda create --name tensorrt python=3.7
conda activate tensorrt
- 官网下载TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz(注意要先下载tensorrt,根据tensorrt的版本来设置cuda和cudnn的版本)
pip install pytorch==1.6.0 torchvision==0.7.0
- cuda和cudnn的安装:conda可以在conda的虚拟环境里安装10.0以上的cudatoolkit和cudnn,但实验发现tensorrt会调用系统中的cuda和cudnn。所以需要在系统环境中安装对应版本的cuda和cudnn。
tar方法安装tensorrt
tar xzvf TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.0.cudnn7.6.tar.gz
vim ~/.bashrc
,添加环境变量,如果使用了zsh,要在 ~/.zshrc中添加该环境变量
export LD_LIBRARY_PATH=/home/ps/Downloads/TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.0.cudnn7.6/TensorRT-7.0.0.11/lib:${LD_LIBRARY_PATH}
cd TensorRT-7.0.0.11/python
pip install tensorrt-7.0.0.11-cp37-none-linux_x86_64.whl
# 注意选择自己python对应版本的whlcd TensorRT-7.0.0.11/graphsurgeon
pip install graphsurgeon-0.4.1-py2.py3-none-any.whl
- 验证tensorrt是否安装成功:
python
import tensorrt
tensorrt.__version__
安装torch2trt(实验失败,可以直接跳到下一步)
- pip安装自己库里涉及的其他库,例如OpenCV,matplotlib等
- 在当前conda环境中,安装torch2trt库
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install
13.验证torch2trt是否安装成功:
import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet
# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])
y = model(x)
y_trt = model_trt(x)
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
以上实验成功,证明tensorrt和torch2trt都安装成功,但是实际使用时,Torch2trt这个库却不能使用成功,捣鼓了好半天后,决定换个工具
TensorRTx
在yolov5转trt的说明中,看到了这个库,这个开源库帮我们写好了很多常见网络的pytorch转tensorrt代码,包括不同版本的yolov5(yolov5官方和tensorrtx官方的版本需要对应)。这个库需要在系统环境下进行编译,所以conda环境无法直接满足需求,编译过程中,容易找不到相关tensorrt库、OpenCV库等,如果要使用这个库,cuda、cudnn和OpenCV还是需要在系统环境中安装。
OpenCV3.4.1版本安装
我本机最后的环境如下:
- Ubuntu18.04
- cuda10.0
- cudnn7.6.5
- python3.7
- pytorch1.7.1 torchvision0.8.2
- tensorrt7.0.0
- OpenCV3.4.1
注意以下几点:
- 在使用该库进行编译时,如果是使用tar方法安装的tensorrt,需要注意修改或者添加CMakeLists.txt里面的tensorrt的include_directories和link_directories为相关路径,否则编译时会找不到tensorrt相关库
- 如果编译的是自己训练的模型,需要在对应c++文件中修改相关输入、输出参数,如image尺寸、模型尺寸等。每次修改后需要重新编译
- 生成的engine文件即为转换成tensorrt格式的模型,可以使用相关api对其进行调用,如果使用python api还需要安装pycuda:
pip install ‘pycuda>=2019.1.1’
可能会遇到Fatal error: cuda.h, No such file or directory
再添加一遍环境变量:export PATH=/usr/local/cuda-10.0/bin:/usr/local/cuda/bin:$PATH