目录
模型导出
在当你训练好你的Yolov8模型时,请务必使用batch=1
选项以确保Yolov8模型能够以batchSize=1
进行输出。
model.export(format="onnx",batch=1,simplify=True)
参数解析
- format:代表你要导出的模型格式,这里我使用了onnx。
- batch : 批次数量
- simplify:是否调用
onnx-simplifier
简化onnx模型
模型输出分解
这里引用一下Ultralytics开发人员glenn-jocher
的英文原文答复来解释一下Yolov8的模型输出:
@RosterMouch(我的Github名) hello,
The output of the YOLOv8 model is a 1D array of length 22xN, where N is the number of detected objects. The first 4 values in the array represent the x, y, width, and height of the bounding box, followed by 18 values representing the confidence scores of each class.
Yolov8的模型输出是一个1维的矩阵,长度为22xN,其中N代表的是模型识别的目标数量。模型的前四位代表着锚点框的x,y,宽度和高度,其余十八个值则代表着每一个类的置信度(在我的任务有18个类,因此18+4=22 x N )。
模型部署
关于如何部署NVIDIA Triton Inference Server,请阅读本篇文章,这里不做赘述
如何设置modelRepository(模型仓库)
Triton要求开发者有严格的模型仓库格式,如下:
└── classnet
├── 1
│ └── model.onnx
└── config.pbtxt
由于这里我仍然没有搞定TensorRT plan模型找不到文件的根本原因,所以这里以onnx模型进行部署。
目录结构解释如下:
- classnet:代表你的模型名称
- 1:代表你的模型版本,如果有大于1的版本时,如果不在
config.pbtxt
文件中指明Triton的加载策略时,Triton默认只会加载最新版本的模型。 - config.pbtxt:模型的配置文件。
如何编写模型配置文件(config.pbtxt):
config.pbtxt文件的常用配置项如下:
name: "classnet" # 模型名称,和你的文件夹一一对应
platform: "onnxruntime_onnx" # 指定模型运行的后端
max_batch_size: 0 # 最大的批次大小,0代表自动
input [ # Triton Input输入对象配置
{
name: "images", # 输入名称(不要动,v8的模型定义就是这个)
data_type: TYPE_FP32, # 根据你的权重类型进行修改,这里我的模型时FP32
dims: [ 1,3,640,640 ] # 输入为1批次,大小为640x640像素的RGB图像
}
]
output[ # Triton Output输出对象配置
{
name: "output0", # 输出名称
data_type: TYPE_FP32 # 输出数据类型
dims: [ 1,22,8400 ] # 输出大小,一般默认是1批次,N个类,8400个目标(当然比这个值小也正常)
}
]
# 版本策略配置
# 其中latest代表Triton加载的最新版本模型
# num_version代表版本号
version_policy: { latest { num_versions: 1 } }
# instance_group:模型运行实例(设备)组配置
instance_group: [
{
count: 1 # 数量
kind: KIND_GPU # 类型
gpus: [ 0 ] # 如果参数项为GPU,则该列表将指定对应序号下的可见CUDA设备来运行模型
}
]
如何编写docker容器启动脚本
关于如何编写docker启动脚本,请阅读Triton部署文章中启动项解析
部分进行参考,这里不做赘述。
客户端
Triton Client安装
在开始执行客户端的脚本编写时,你需要安装Triton Client
Python包,安装方法如下:
pip install tritonclient[all] -i https://pypi.tuna.tsinghua.edu.cn/simple
核心代码编写
客户端构建
首先我们要从tritonclinet
中引入http
模块,以及针对处理推理请求可能时发生的异常类,通过引入tritonclient.utils.InferenceServerException
:
import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException
之后我们需要创建一个随机的哈希来用作表示推理请求的ID:
# Creat a random hash for request ID
requestID = random.randint(0, 100000)
requestID = sha256(str(requestID).encode('utf-8')).hexdigest()
接下来实例化Triton客户端:
# Create a HTTP client for inference server running in localhost:8000.
triton_client = httpclient.InferenceServerClient(
url="localhost:8000",
)
并未为模型的输入和输出创建列表。
前处理
请切记,在将opencv图像流转换为numpy矩阵后,一定要为其扩增维度,因为我们的模型接受的是:(1,3,640,640)
的图像,对于此,我们可以使用 numpy.expand_dims()
函数对其进行维度扩增:
imageData = cv.imread("./original.jpg")
imageData = np.array(imageData)
imageData = np.expand_dims(imageData, axis=0)
# resize into 640
imageData = np.resize(imageData, (1,3,640,640))
# change to fp32
imageData = imageData.astype(np.float32)
对象初始化
接下来创建一个推理请求:
inputs.append(httpclient.InferInput('images', imageData.shape, "FP32"))
其中第一位对应的是你模型的输入名称,第二位对应的是输入的大小(1,3,640,640),第三位代表数据类型。
但这时模型的输入并不完整,我们只是初始化了一个推理请求,但并没有数据,所以我们可以调用httpclient.InferInput.set_data_from_numpy()
方法为其添加数据:
inputs[0].set_data_from_numpy(imageData)
与之对应,我们也要创建一个模型请求后的输出:
outputs.append(httpclient.InferRequestedOutput('output0'))
其中output0
代表模型的输出名称。
执行推理请求
接下来,我们开始执行推理请求了:调用httpclient.async_infer()
函数执行异步推理请求:
try:
resp = triton_client.async_infer(
model_name="classnet", # 模型名称
model_version="1", # 版本号
inputs=inputs, # 输入
outputs=outputs, # 输出
request_id=requestID # 请求ID
)
except InferenceServerException as e:
print(e)
这样,我们就像Triton服务器发送了一个推理请求。
后处理
当我们得到了正确的响应后,我们应该怎么办呢?
我们知道,模型的输出将是一个1D的矩阵,其中N是类别数量,M是识别的目标数量,但是这显然对于后续的如:结果提取和锚点的绘制。
因此我们需要将这一结果进行转置:
result = resp.get_result().as_numpy('output0')[0].T
其中.get_result()
函数可以从响应中获取结果,并使用.as_numpy()
分离矩阵结果,output0
代表模型的输出名,
这样我们就拿到了结果。
还记得我们对模型输出的定义吗?基于此,我们得到了一个后处理函数:
@numba.jit
def extractFeature(inferenceResult:np.ndarray) -> np.ndarray:
threshold = 0.5 # 设置置信度阈值
final = [] # 创建最终输出的数组
for i in inferenceResult: # 遍历结果
temp = [] # 创建临时结果数组
box = i[:4] # 截取锚点框信息
confidence = i[4:] # 截取所有类别的置信度
if np.max(confidence) < threshold:
continue # 如果当前置信度小于阈值,则抛弃推理结果
else:
for j in box: # 分解锚点框信息
temp.append(j)
# 获取最佳置信度对应的类别
classIndex = np.where(confidence == np.max(confidence))[0][0]
temp.append(classIndex)
final.append(temp)
final = np.array(final)
# drop duplicate boxes
# final = np.unique(final, axis=1)
return final
接下来,就是如何将[x,y,w,h]
类型的数据转换为[x1,y,1,x2,y2]
格式的数据了,因为Yolo的默认输出不适用于OpenCV的rectangle()函数,因此我们需要对数据进行一次转换:
@numba.jit
def xywh2xyxy(box:np.ndarray) -> np.ndarray:
x = box[0]
y = box[1]
w = box[2]
h = box[3]
x1 = x - w/2 # 将原始的x - 二分之一的宽度
y1 = y - h/2 # 将原始的y - 二分之以的高度
x2 = x + w/2 # 同理将x2和y2与x1,y1做反方向操作即可
y2 = y + h/2
return np.array([x1,y1,x2,y2])
接下来就是渲染函数了:
@numba.jit
def drawBox(image:np.ndarray, result:np.ndarray) -> np.ndarray:
for i in result:
x = int(i[0])
y = int(i[1])
classIndex = str(int(i[4]))
boxWidth = int(i[2]) - int(i[0])
boxHeight = int(i[3]) - int(i[1])
# random RGB
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
cv.rectangle(image, (x, y), (x+boxWidth, y+boxHeight),color, 2)
cv.putText(image,str(classIndex), (x, y), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
return image
完整代码
import numpy as np
import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException
import random
from hashlib import sha256
import numba
import cv2 as cv
import json
classRGB = json.loads(open("../config/RGBColorDict.json").read())
# keys to list
className = list(classRGB.keys())
classColors = list(classRGB.values())
@numba.jit
def extractFeature(inferenceResult:np.ndarray) -> np.ndarray:
threshold = 0.95
final = []
for i in inferenceResult:
temp = []
box = i[:4]
confidence = i[4:]
if np.max(confidence) < threshold:
continue
else:
for j in box:
temp.append(j)
classIndex = np.where(confidence == np.max(confidence))[0][0]
temp.append(classIndex)
final.append(temp)
final = np.array(final)
# drop duplicate boxes
# final = np.unique(final, axis=1)
return final
@numba.jit
def xywh2xyxy(box:np.ndarray) -> np.ndarray:
x = box[0]
y = box[1]
w = box[2]
h = box[3]
x1 = x - w/2
y1 = y - h/2
x2 = x + w/2
y2 = y + h/2
return np.array([x1,y1,x2,y2])
@numba.jit
def drawBox(image:np.ndarray, result:np.ndarray) -> np.ndarray:
for i in result:
# convet [x,y,w,h] to [x1,y1,x2,y2]
box = xywh2xyxy(i[:4])
classIndex = className[int(i[4])]
classColor = classColors[int(i[4])]
# random RGB
color = (classColor[0], classColor[1], classColor[2])
# draw box
cv.rectangle(image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 2)
# draw text
cv.putText(image, classIndex, (int(box[0]), int(box[1])), cv.FONT_HERSHEY_SIMPLEX, 1, color, 2)
return image
def main():
# Creat a random hash for request ID
requestID = random.randint(0, 100000)
requestID = sha256(str(requestID).encode('utf-8')).hexdigest()
# Create a HTTP client for inference server running in localhost:8000.
triton_client = httpclient.InferenceServerClient(
url="localhost:8000",
)
inputs = []
outputs = []
imageData = cv.imread("./original.jpg")
imageData = np.array(imageData)
imageData = np.expand_dims(imageData, axis=0)
# resize into 640
imageData = np.resize(imageData, (1,3,640,640))
# change to fp32
imageData = imageData.astype(np.float32)
inputs.append(httpclient.InferInput('images', imageData.shape, "FP32"))
inputs[0].set_data_from_numpy(imageData)
outputs.append(httpclient.InferRequestedOutput('output0'))
# Send request to the inference server. Get results for both output tensors.
try:
resp = triton_client.async_infer(
model_name="classnet",
model_version="1",
inputs=inputs,
outputs=outputs,
request_id=requestID
)
result = resp.get_result().as_numpy('output0')[0].T
result = np.squeeze(result)
result = extractFeature(result)
image = cv.imread("./original.jpg")
image = cv.resize(image, (640, 640))
imageOutput = drawBox(image, result)
cv.imwrite("output.jpg", imageOutput)
except InferenceServerException as e:
print(e)
if __name__ == '__main__':
main()