triton 部署yolov5模型

一、安装triton

0.安装docker和nvidia-docker

1. 安装以下网站获取Nvidia授权账号密码

Nvidia授权

2. 安装triton

docker login nvcr.io

Username: XXXXXXX
Password: XXXXXXX

sudo docker pull nvcr.io/nvidia/tritonserver:22.04-py3

二、准备yolov5模型

参考 yolov5训练自己的数据集

三、启动triton服务

1.搭建模型目录

/home/triton/model_repository  # model-repository-path 可自行修改 
└── model_test  # your_model_name  可自行修改
      ├── 1  # 固定路径，不可修改
      │   └── model.onnx    # 默认onnx模型名

2.启动服务

sudo docker run -d --gpus=1 --rm -p8000:8000 -p8001:8001 -p8002:8002 -v /home/triton/model_repository:/models nvcr.io/nvidia/tritonserver:22.04-py3 tritonserver --model-repository=/models --strict-model-config=false

4. 查看config，登录网址

http://ip:8000/v2/models/model_test/config

三、客户端调用

1.需要tritonclient模块，参考代码

import numpy as np
import tritonclient.grpc as grpcclient
import cv2 as cv

color_list = [(255, 0, 0), (0, 255, 0)]

class MY_data:
    def __init__(self, rect_list, name_list, img) -> None:
        self.rect_list = rect_list
        self.name_list = name_list
        self.img = img

    def show(self):
        for i in range(len(self.rect_list)):
            cv.rectangle(self.img, self.rect_list[i], color_list[0] if self.name_list[i] == 'people' else color_list[1])
        cv.imshow("test", self.img)
        cv.waitKey(0)
        cv.imwrite('test.jpg', self.img)

if __name__ == "__main__":

    triton_client =  grpcclient.InferenceServerClient(
        url='XXX.XXX.XXX.XXX:8001',
        verbose=False,
        ssl=False,
        root_certificates=None,
        private_key=None,
        certificate_chain=None)
    # 测试图片    
    src_img = cv.imread('3.jpg')

    max_side = src_img.shape[1] if src_img.shape[1] >src_img.shape[0] else src_img.shape[0]

    dst_img = np.zeros([max_side, max_side, 3], dtype=np.uint8)
    dst_img[0:src_img.shape[0], 0:src_img.shape[1]] = src_img.copy()


    x_factor = dst_img.shape[1]/640
    y_factor = dst_img.shape[0]/640

    img = cv.resize(dst_img, [640, 640])
    
    img = img /255.
    inputs = []
    inputs.append(grpcclient.InferInput('images', [3, 640, 640], "FP32"))
    img = img.astype(np.float32).transpose((2, 0, 1))

    # 增加维度，编程 1X3X640X640
    # img = np.expand_dims(img,axis=0)

    inputs[0].set_data_from_numpy(img)
    outputs = []
    outputs.append(grpcclient.InferRequestedOutput('output'))

    results = triton_client.infer(model_name='yolov5', inputs=inputs, outputs=outputs,compression_algorithm=None)

    predictions = results.as_numpy("output")[0]
    # print(predictions[0])

    rows = 25200
    dimensions = 7

    scoreThreshold = 0.2
    nmsThreshold = 0.4
    confThreshold = 0.4

    class_name = ['label1', 'label2']

    confidence_list = []
    class_list = []
    rect_list = []

    for i in range(rows):
        res = predictions[i]
        confidence = res[4]
        
        if confidence > confThreshold:
            print(confidence)
            max = -1
            max_index = -1
            for j in range(len(class_name)):
                if res[5 + j] > max:
                    max = res[5 + j]
                    max_index = j
            confidence_list.append(confidence)
            class_list.append(class_name[max_index])
            x = res[0]
            y = res[1]
            w = res[2]
            h = res[3]
            left = int((x - w/2)*x_factor)
            top = int((y - h/2)*y_factor)
            width = int(w*x_factor)
            height = int(h*y_factor)
            rect_list.append([left, top, width, height])
    # print(rect_list)
    nms_result = cv.dnn.NMSBoxes(np.array(rect_list), np.array(confidence_list), scoreThreshold, nmsThreshold)  
    print(nms_result)  
    r_l = []
    n_l = []
    for k in range(len(nms_result)):
        id = nms_result[k]
        rect = rect_list[id]
        r_l.append(rect)
        name = class_list[id]
        n_l.append(name)
    # print(r_l)
    my_data = MY_data(r_l, n_l, src_img)
    my_data.show()