NVIDIA Triton Inference Server 部署 yolov5

一、拉取 tensorrt 、yolov5、tritonserver 镜像

docker pull hakuyyf/tensorrtx:trt8.2_cuda11.2
docker pull ultralytics/yolov5:v6.2
docker pull nvcr.io/nvidia/tritonserver:22.05-py3 

需要注意的是， tensorrtx 镜像和 triton 镜像中的 tensorrt 版本要保持一致， 否则 triton 无法正常调用模型，triton 镜像中的 tensorrt 版本可通过这里查询
Triton Inference Server Release 22.05

二、下载 yolov5-6.2、tensorrtx/yolov5-6.2源码

mkdir triton-yolov5-trt
cd triton-yolov5-trt# 下载 yolov5-6.2 源码
git clone -b v6.2 https://github.com/ultralytics/yolov5.git# 下载 tensorrtx/yolov5-6.2 源码
git clone -b yolov5-v6.2 https://github.com/wang-xinyu/tensorrtx.git

三、pt转wts

开一个新终端

cd triton-yolov5-trt
cp tensorrtx/yolov5/gen_wts.py yolov5
cd yolov5# 进入yolov5-6.2镜像
docker run -it --rm --gpus all -v $PWD:/yolov5 ultralytics/yolov5:v6.2 /bin/bash  # 在容器内操作, yolov5s.pt需要6.2版本的
cd /yolov5
python gen_wts.py -w yolov5s.pt -o yolov5s.wts

四、wts转engine

开一个新终端

cd triton-yolov5-trt
cp yolov5/yolov5s.wts tensorrtx/yolov5
cd tensorrtx/yolov5  
docker run -it --rm --gpus all -v $PWD:/yolov5 hakuyyf/tensorrtx:trt8.2_cuda11.2 /bin/bash   
cd /yolov5
mkdir build  
cd build   
cmake ..  
make -j16  
./yolov5 -s ../yolov5s.wts ../yolov5s.engine s 

执行完将得到 yolov5s.engine 和 libmyplugins.so

五、创建triton推理服务器

开一个新终端

cd triton-yolov5-trt
mkdir -p triton_deploy/models/yolov5/1/  
mkdir triton_deploy/plugins  
cp tensorrtx/yolov5/yolov5s.engine triton_deploy/models/yolov5/1/model.plan  
cp tensorrtx/yolov5/build/libmyplugins.so triton_deploy/plugins/libmyplugins.so  

启动triton

docker run \
--gpus all \
--rm \
-p9000:8000 -p9001:8001 -p9002:8002 \
-v $(pwd)/triton_deploy/models:/models \
-v $(pwd)/triton_deploy/plugins:/plugins \
--env LD_PRELOAD=/plugins/libmyplugins.so \
nvcr.io/nvidia/tritonserver:22.05-py3 tritonserver \
--model-repository=/models \
--strict-model-config=false \
--log-verbose 1

六、创建客户端进行测试

开一个新终端

cd triton-yolov5-trt
git clone https://github.com/tienluongngoc/yolov5_triton_inference_server.git
cd yolov5_triton_inference_server/client/yolov5
pip install tritonclient
python client.py -o data/dog_result.jpg image data/dog.jpg