文章目录
- 一、拉取 tensorrt 、yolov5、tritonserver 镜像
- 二、下载 yolov5-6.2、tensorrtx/yolov5-6.2源码
- 三、pt转wts
- 四、wts转engine
- 五、创建triton推理服务器
- 六、创建客户端进行测试
一、拉取 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