yolov4 + pytorch + tensorrt + cuda + docker环境配置
Published:
yolov4 + pytorch + tensorrt + cuda + docker环境配置。各主要软件版本如下：
pytorch == 1.6.0
torchvision == 0.7.0
cuda == 10.2
cudnn == 7.6
tensorrt == 7.0
os == Ubuntu18.04
使用的yolo源代码为bubbliiiing/yolov4-tiny-pytorch。将pytorch模型转换为tensorrt模型用了两种方式，方式一为使用NVIDIA的NVIDIA-AI-IOT/torch2trt。方式二为将pytorch模型转换为ONNX，然后转化为tensorrt模型。
首先创建文件夹tensorrt_docker，将bubbliiiing/yolov4-tiny-pytorch和NVIDIA-AI-IOT/torch2trt克隆到文件夹中。从链接：https://pan.baidu.com/s/10nnhtur_6yz5wY4ksTfQsQ 提取码：71op，下载tensorrt的安装压缩包文件。tensorrt安装使用whl的方式安装，没有采用deb方式安装是因为deb方式安装要求cuda以及cudnn同样是使用deb方式安装，使用nvidia提供的cuda环境的情况下，非常容易出现依赖问题。将tensorrt的压缩包文件同样放到tensorrt_docker文件夹中。
在tensorrt_docker文件夹中创建Dockerfile文件。
Dockerfile
#使用nvidia提供的cuda和cudnn环境作为基础镜像，devel分支可以正常使用nvcc等指令，一些依赖的编译需要nvcc。
FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04
#此命令可以指令镜像构建过程中使用的shell，默认是sh，因此改为常用的bash。
SHELL ["/bin/bash","-c"]
#ffmpeg、libsm6、libxext6为opencv-python的依赖文件
RUN apt update && apt install -y python3 \
python3-pip \
swig \
git \
vim \
ffmpeg \
libsm6 \
libxext6 \
&& rm -rf /var/lib/apt/lists/*
#更新pip
RUN pip3 install -U pip
WORKDIR /
#使用ADD指令会自动进行解压，解压后的文件夹名称为TensorRT-7.0.0.11，位于container内的 / 目录下。
ADD TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn7.6.tar.gz /
#安装tensorrt需要
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/TensorRT-7.0.0.11/lib:/TensorRT-7.0.0.11/targets/x86_64-linux-gnu/lib:/usr/local/cuda/lib64" >> /root/.bashrc && source /root/.bashrc
#ENV指令将环境变量暴漏在构建镜像的过程中，与上面的环境变量的设置不同，上面的环境变量内容将会被写入.bashrc中。镜像构建过程中需要使用ENV显式指定环境变量。
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/TensorRT-7.0.0.11/lib:/TensorRT-7.0.0.11/targets/x86_64-linux-gnu/lib:/usr/local/cuda/lib64
#使用对应的whl文件安装tensorrt
RUN cd /TensorRT-7.0.0.11/python && python3 -m pip install tensorrt-7.0.0.11-cp36-none-linux_x86_64.whl
RUN cd /TensorRT-7.0.0.11/graphsurgeon && python3 -m pip install graphsurgeon-0.4.1-py2.py3-none-any.whl
#安装torch、torchvision、opencv、matplotlib等文件，具体见requirements.txt
COPY ./requirements.txt /
RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
RUN pip install -r /requirements.txt
#安装torch2trt，此lib的作用为将模型直接转为tensorrt模型
COPY ./torch2trt /torch2trt
RUN cd /torch2trt && python3 setup.py install
在tensorrt_docker文件夹中创建requirements.txt文件。
requirements.txt
torch==1.6.0
torchvision==0.7.0
tensorboard
matplotlib
tqdm
packaging
scipy
opencv-python
执行
sudo docker image build -t tensorrt:torch .
构建镜像。