YOLOV3在Triton Inference Server部署并进行对象检测

注：原 TensorRT Inference Server 官方已改名为 Triton Inference Server

需要的镜像文件

• nvcr.io/nvidia/tensorrtserver:19.10-py3
• nvcr.io/nvidia/tensorrt:19.10-py2
• enisberk/tensorrtserver_client:19.05

1、生成YOLOV3的.plan模型

首先在 YOLO官网 下载YOLOV3的yolov3.weights

Clone代码 yolov3-tensorrt

启动tensorrt容器：

docker run -it --name trt -v /yolov3-tensorrt:/yolov3-tensorrt nvcr.io/nvidia/tensorrt:19.10-py2

进入容器并构建yolov3.plan：

# 进入容器（之后步骤都是在容器内执行）
docker exec -it trt /bin/bash
# 进入yolov3-tensorrt目录，执行yolov3_to_onnx脚本。
cd /yolov3-tensorrt
# 此步需要自行下载python依赖，可以根据执行后的import Error下载
python yolov3_to_onnx.py
# 编译trtexec
cd /workspace/tensorrt/samples/trtexec && make
# 转换文件
cd /workspace/tensorrt/
./bin/trtexec --onnx=$YOUR_PATH/yolov3.onnx --saveEngine=$YOUR_PATH/model.plan

PS：新版本的tenosrrt中，trtexec的参数不是saveEngine，输入trtexec -h查看对应命令

2、将模型部署到Triton Inference Server中

编写YOLOV3模型的config.pbtxt：

name: "yolov3"
platform: "tensorrt_plan"
max_batch_size: 1
default_model_filename: "model.plan"
input [
  {
    name: "000_net"
    data_type: TYPE_FP32
    dims: [ 3, 608, 608 ]
    format: FORMAT_NCHW
  }
]
output [
    {
    name: "082_convolutional"
    data_type: TYPE_FP32
    dims: [ 255, 19, 19]
    },
    {
    name: "094_convolutional"
    data_type: TYPE_FP32
    dims: [ 255, 38, 38 ]
    },
    {
    name: "106_convolutional"
    data_type: TYPE_FP32
    dims: [ 255, 76, 76 ]
    }
]
instance_group [
  {
    count:1
    kind: KIND_GPU
  }
]

然后部署到Triton Inference Server中

模型文件结构：

yolov3
├── 1
│   └── model.plan
└── config.pbtxt

启动server容器：

# 将上面yolov3模型整个目录放置在你的模型目录下
nvidia-docker run -itd --name trt_server --shm-size=1g \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -p8000:8000 -p8001:8001 -p8002:8002 \
  -v /你的模型目录:/models \
  nvcr.io/nvidia/tensorrtserver:19.10-py3 trtserver --model-store=/models

查看模型部署状态：

curl localhost:8000/api/status

3、使用yolo_client.py进行检测

先启动tensorRT client容器并进入容器：

# 启动client容器
docker run -it --name trt_client -v /yolov3-tensorrt:/yolov3-tensorrt enisberk/tensorrtserver_client:19.05
# 进入client容器
docker exec -it trt_client /bin/bash

执行yolo_client.py脚本：

# 先切换目录
cd /yolov3-tensorrt
# 各项配置在脚本中均可找到
python yolo_client.py -m yolov3

检测结果图片会存放到./detected_image目录下。