考虑实际落地中部署推理资源消耗,速度与模型精度之间的均衡,同时,车辆属于比较通用场景的目标检测,我们将YOLOV11作为最终的模型选择,目前官方开源的模型基于coco数据集训练,里面已经包含了车辆相关类别,经测试,在监控视角下效果基本可用,所以,在落地过程中,我们先用这个模型来打通部署链路,后面再逐步优化模型性能。
在Nvidia的硬件下做部署,部署框架优先选择官方框架TensorRT,下面重点讲述如何用TensorRT部署YOLOV11模型。主要分为:训练框架下模型的推理验证,中间件模型onnx的转换,中间件到TensorRT模型转换,以及C++部署四个部分,每一步都会穿插着精度对齐,最后会做推理性能测试。
训练框架下模型的推理验证
一般思路下讲一个模型的应用会按模型理论知识,训练优化,推理部署这样的思路来。我们以项目推进角度整个流程就反过来了,先讲部署,项目上急着用呢,模型的部署要关注三个部分,即模型的前处理,模型的推理以及模型的后处理部分,一般会在推理框架上先做验证,确认没问题后再向后面的部署框架迁移。这三部分我们最关注的是前处理和后处理,推理更多的是框架自己搞定。
模型前处理
YOLOV11的默认模型前处理过程为不变形resize,默认尺寸为640*640,即:
-
计算缩放比例,原始图像宽大于高,则将宽设置为640,高度自适应到相应的值;原始图像宽小于高,则高设置为640,宽度自适应到响应的值;
-
使用opencv自带的resize函数将图片按上述计算缩放尺寸进行缩放,方式选择“双线性差值”。然后再构建一个640*640的底图,像素值均初始化为(114,114,114),并将已缩放好的原始图片居中贴在底图上。
-
将opencv的BGR格式图片数据转换为RGB,并进行归一化操作,并转换为维度格式为[N,C,H,W]四维向量输入模型,本期推理默认batch为1。
具体代码参考:
def preprocess_warpAffine(image, dst_width=640, dst_height=640):
"""
图像预处理函数,用于将输入图像变形并缩放到指定大小。
该函数首先计算缩放因子,以确保图像按比例缩放到目标尺寸,
然后计算偏移量,以确保缩放后的图像居中。使用计算得到的缩放
和偏移参数对图像进行仿射变换,并对变换后的图像进行颜色反转、
归一化和维度转换,以适应深度学习模型的输入要求。
参数:
image: 输入的图像,通常是一个二维或三维的numpy数组。
dst_width: 目标图像的宽度,默认值为640像素。
dst_height: 目标图像的高度,默认值为640像素。
返回:
img_pre: 预处理后的图像,是一个四维的PyTorch张量。
IM: 逆仿射变换矩阵,用于将变形后的图像恢复到原始状态。
"""
# 计算缩放因子,确保图像按比例缩放到目标尺寸
scale = min((dst_width / image.shape[1], dst_height / image.shape[0]))
# 计算偏移量,确保缩放后的图像居中
ox = int((dst_width - scale * image.shape[1]) / 2)
oy = int((dst_height - scale * image.shape[0]) / 2)
# 构建仿射变换矩阵
M = np.array([
[scale, 0, ox],
[0, scale, oy]
], dtype=np.float32)
# 对图像进行仿射变换,设置目标尺寸、插值方法和边界填充方式
img_pre = cv2.warpAffine(image, M, (dst_width, dst_height), flags=cv2.INTER_LINEAR,
borderMode=cv2.BORDER_CONSTANT, borderValue=(114, 114, 114))
# 计算逆仿射变换矩阵,用于将变形后的图像恢复到原始状态
IM = cv2.invertAffineTransform(M)
# 对变换后的图像进行颜色反转和归一化处理
img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32)
# 调整图像数据的维度顺序,适应深度学习模型的输入要求
img_pre = img_pre.transpose(2, 0, 1)[None]
# 将处理后的图像数据转换为PyTorch张量
img_pre = torch.from_numpy(img_pre)
# 设置PyTorch打印选项,提高输出精度
torch.set_printoptions(precision=6)
# 返回预处理后的图像和逆仿射变换矩阵
return img_pre, IM
模型后处理
模型经推理后,输出格式为[1, 84, 8400],其中84表示[cx,cy,w,h,class×80],即检测框的中心点坐标,宽,高,以及coco数据集80个类别的置信度;8400表示三个不同尺度的输出,即20×20+40×40+80×80。
-
按设置的置信度阈值过滤8400个候选框,默认0.25;
-
对于过滤后的候选框执行NMS,iou阈值默认0.45,剩下的即为输出检测框。
具体代码参考:
# 计算两个边界框的交并比(Intersection over Union,IoU)
def iou(box1, box2):
# 内部函数,用于计算边界框的面积
def area_box(box):
return (box[2] - box[0]) * (box[3] - box[1])
# 计算交集的左上角和右下角坐标
left = max(box1[0], box2[0])
top = max(box1[1], box2[1])
right = min(box1[2], box2[2])
bottom = min(box1[3], box2[3])
# 计算交集的面积
cross = max((right-left), 0) * max((bottom-top), 0)
# 计算并集的面积
union = area_box(box1) + area_box(box2) - cross
# 如果交集或并集的面积为0,则返回0
if cross == 0 or union == 0:
return 0
# 返回交并比
return cross / union
# 非极大值抑制(Non-Maximum Suppression,NMS)算法,用于去除多余的边界框
def NMS(boxes, iou_thres):
# 初始化一个标记列表,用于标记是否移除某个边界框
remove_flags = [False] * len(boxes)
# 初始化一个列表,用于保存保留下来的边界框
keep_boxes = []
# 遍历所有边界框
for i, ibox in enumerate(boxes):
# 如果当前边界框已被标记为移除,则跳过
if remove_flags[i]:
continue
# 将当前边界框添加到保留列表中
keep_boxes.append(ibox)
# 继续遍历其余的边界框
for j in range(i + 1, len(boxes)):
# 如果当前边界框已被标记为移除,则跳过
if remove_flags[j]:
continue
# 获取下一个边界框
jbox = boxes[j]
# 如果两个边界框的类别标签不同,则跳过
if(ibox[5] != jbox[5]):
continue
# 如果两个边界框的交并比大于设定的阈值,则标记第二个边界框为移除
if iou(ibox, jbox) > iou_thres:
remove_flags[j] = True
# 返回保留下来的边界框列表
return keep_boxes
# 后处理函数,用于处理模型的预测结果
def postprocess(pred, IM=[], conf_thres=0.25, iou_thres=0.45):
# 初始化一个列表,用于保存处理后的边界框
boxes = []
# 遍历模型的预测结果
for item in pred[0]:
# 提取边界框的中心坐标、宽度、高度
cx, cy, w, h = item[:4]
# 提取类别标签和置信度
label = item[4:].argmax()
confidence = item[4 + label]
# 如果置信度低于设定的阈值,则跳过
if confidence < conf_thres:
continue
# 计算边界框的左上角和右下角坐标
left = cx - w * 0.5
top = cy - h * 0.5
right = cx + w * 0.5
bottom = cy + h * 0.5
# 将边界框的坐标、置信度和类别标签添加到列表中
boxes.append([left, top, right, bottom, confidence, label])
# 将边界框列表转换为NumPy数组
boxes = np.array(boxes)
# 如果没有边界框,则返回空列表
if 0 == len(boxes):
return []
# 应用图像的缩放比例和偏移量,将边界框坐标转换回原始图像的坐标系
lr = boxes[:,[0, 2]]
tb = boxes[:,[1, 3]]
boxes[:,[0,2]] = IM[0][0] * lr + IM[0][2]
boxes[:,[1,3]] = IM[1][1] * tb + IM[1][2]
# 根据置信度降序排序边界框
boxes = sorted(boxes.tolist(), key=lambda x:x[4], reverse=True)
# 应用非极大值抑制,去除多余的边界框
return NMS(boxes, iou_thres)
数据集下的推理验证
上述做完后,需要做一下测试集上的推理验证,我们选择coco数据集自带的5000张测试集,验证下来结果基本没有问题。
同时,按“图片名称,类别索引,置信度,检测框”格式,保存识别结果,用于后面模型转换精度保持的精度分析。
中间件模型onnx的转换
本次使用pytorch转TensorRT的常规方式,通过转换中间件onnx进行转换,这步还比较简单,不BB,直接见代码:
model = YOLO(model='yolo11x.pt') # load a pretrained model (recommended for training)
model.export(format="onnx", opset=16, imgsz=640, simplify=True, save_dir='.')
onnx的前后处理也pytorch一样,只是推理调用部分有所区别,这部分不放代码了,后面大家自己在代码仓库中看即可。
中间件到TensorRT模型转换
中间件模型转换完以后,正式到转TensorRT模型的阶段了,我们这次采用的是调用TensorRT的python接口的方式转换(方式有很多,选合适的即可),为快速跑通链路,我们选择了FP16精度做部署,后面细化调优的时候,会跟大家讲讲INT8,这里面可能遇到的坑比较多,到时候细聊。同样不BB,直接上代码:
def build_engine(onnx_model_path, engine_path, precision_set="fp16"):
"""
从ONNX模型生成TensorRT引擎。
参数:
onnx_model_path: str,ONNX模型的路径。
engine_path: str,生成的TensorRT引擎的保存路径。
precision_set: str,精度设置,默认为"fp16"。
返回:
无
"""
# 创建TensorRT日志记录器
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# 创建TensorRT构建器
builder = trt.Builder(TRT_LOGGER)
# 创建网络定义,并设置显式批处理标志
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 创建ONNX解析器
parser = trt.OnnxParser(network, TRT_LOGGER)
# 打开ONNX模型文件并解析
with open(onnx_model_path, "rb") as model:
if not parser.parse(model.read()):
# 如果解析失败,打印错误信息
for error in range(parser.num_errors):
print(parser.get_error(error))
# 创建构建配置
config = builder.create_builder_config()
# 设置最大工作空间大小为1GB
config.max_workspace_size = 1 << 30 # 1GB
# 如果精度设置为fp16,启用FP16模式
if precision_set == "fp16":
config.set_flag(trt.BuilderFlag.FP16)
# 如果引擎文件不存在,构建并保存引擎
if not os.path.exists(engine_path):
engine = builder.build_engine(network, config)
with open(engine_path, "wb") as f:
# 序列化引擎并写入文件
f.write(engine.serialize())
模型转换完,也要做个这部分的单元测试,TensorRT模型的前后处理和pytorch模型不太一样,不过差异不大,主要是:
-
预处理阶段转成[N,C,H,W]四维向量后,需要拉成一维传给模型;
-
模型推理完成以后,也是以一维的方式输出,同样要转换为[1, 84, 8400]格式后,接前面的后处理算法。
具体代码不放了,差异不大,见仓库。
C++部署
最终算法将以C++的方式部署在客户机器上,所以需要将TensorRT模型python部分的推理代码移植成c++,同时顺带做了些推理优化,主要是将前处理和后处理的代码用cuda重写了,以提升推理性能。
前处理的cuda移植
-
双线性差值的不变形resize,做灰边填充部分的代码如下:
__global__ void letterbox(const uchar* srcData, const int srcH, const int srcW, uchar* tgtData,
const int tgtH, const int tgtW, const int rszH, const int rszW, const int startY, const int startX)
{
int ix = threadIdx.x + blockDim.x * blockIdx.x;
int iy = threadIdx.y + blockDim.y * blockIdx.y;
int idx = ix + iy * tgtW;
int idx3 = idx * 3;
if ( ix > tgtW || iy > tgtH ) return; // thread out of target range
// gray region on target image
if ( iy < startY || iy > (startY + rszH - 1) ) {
tgtData[idx3] = 114;
tgtData[idx3 + 1] = 114;
tgtData[idx3 + 2] = 114;
return;
}
if ( ix < startX || ix > (startX + rszW - 1) ){
tgtData[idx3] = 114;
tgtData[idx3 + 1] = 114;
tgtData[idx3 + 2] = 114;
return;
}
float scaleY = (float)rszH / (float)srcH;
float scaleX = (float)rszW / (float)srcW;
// (ix,iy)为目标图像坐标
// (before_x,before_y)原图坐标
float beforeX = float(ix - startX + 0.5) / scaleX - 0.5;
float beforeY = float(iy - startY + 0.5) / scaleY - 0.5;
// 原图像坐标四个相邻点
// 获得变换前最近的四个顶点,取整
int topY = static_cast<int>(beforeY);
int bottomY = topY + 1;
int leftX = static_cast<int>(beforeX);
int rightX = leftX + 1;
//计算变换前坐标的小数部分
float u = beforeX - leftX;
float v = beforeY - topY;
if (topY >= srcH - 1 && leftX >= srcW - 1) //右下角
{
for (int k = 0; k < 3; k++)
{
tgtData[idx3 + k] = (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k];
}
}
else if (topY >= srcH - 1) // 最后一行
{
for (int k = 0; k < 3; k++)
{
tgtData[idx3 + k]
= (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]
+ (u) * (1. - v) * srcData[(rightX + topY * srcW) * 3 + k];
}
}
else if (leftX >= srcW - 1) // 最后一列
{
for (int k = 0; k < 3; k++)
{
tgtData[idx3 + k]
= (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]
+ (1. - u) * (v) * srcData[(leftX + bottomY * srcW) * 3 + k];
}
}
else // 非最后一行或最后一列情况
{
for (int k = 0; k < 3; k++)
{
tgtData[idx3 + k]
= (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]
+ (u) * (1. - v) * srcData[(rightX + topY * srcW) * 3 + k]
+ (1. - u) * (v) * srcData[(leftX + bottomY * srcW) * 3 + k]
+ u * v * srcData[(rightX + bottomY * srcW) * 3 + k];
}
}
}
-
颜色通道转换以及归一化部分的代码如下:
__global__ void process(const uchar* srcData, float* tgtData, const int h, const int w)
{
int ix = threadIdx.x + blockIdx.x * blockDim.x;
int iy = threadIdx.y + blockIdx.y * blockDim.y;
int idx = ix + iy * w;
int idx3 = idx * 3;
if (ix < w && iy < h)
{
tgtData[idx] = (float)srcData[idx3 + 2] / 255.0; // R pixel
tgtData[idx + h * w] = (float)srcData[idx3 + 1] / 255.0; // G pixel
tgtData[idx + h * w * 2] = (float)srcData[idx3] / 255.0; // B pixel
}
}
后处理的cuda移植
-
数据解码的cuda化
// ------------------ decode ( get class and conf ) --------------------
__global__ void decode_kernel(float* src, float* dst, int numBboxes, int numClasses, float confThresh, int maxObjects, int numBoxElement){
int position = blockDim.x * blockIdx.x + threadIdx.x;
if (position >= numBboxes) return;
float* pitem = src + (4 + numClasses) * position;
float* classConf = pitem + 4;
float confidence = 0;
int label = 0;
for (int i = 0; i < numClasses; i++){
if (classConf[i] > confidence){
confidence = classConf[i];
label = i;
}
}
if (confidence < confThresh) return;
int index = (int)atomicAdd(dst, 1);
if (index >= maxObjects) return;
float cx = pitem[0];
float cy = pitem[1];
float width = pitem[2];
float height = pitem[3];
float left = cx - width * 0.5f;
float top = cy - height * 0.5f;
float right = cx + width * 0.5f;
float bottom = cy + height * 0.5f;
float* pout_item = dst + 1 + index * numBoxElement;
pout_item[0] = left;
pout_item[1] = top;
pout_item[2] = right;
pout_item[3] = bottom;
pout_item[4] = confidence;
pout_item[5] = label;
pout_item[6] = 1; // 1 = keep, 0 = ignore
}
void decode(float* src, float* dst, int numBboxes, int numClasses, float confThresh, int maxObjects, int numBoxElement, cudaStream_t stream){
cudaMemsetAsync(dst, 0, sizeof(int), stream);
int blockSize = 256;
int gridSize = (numBboxes + blockSize - 1) / blockSize;
decode_kernel<<<gridSize, blockSize, 0, stream>>>(src, dst, numBboxes, numClasses, confThresh, maxObjects, numBoxElement);
}
-
nms的cuda化
__device__ float box_iou(
float aleft, float atop, float aright, float abottom,
float bleft, float btop, float bright, float bbottom
){
float cleft = max(aleft, bleft);
float ctop = max(atop, btop);
float cright = min(aright, bright);
float cbottom = min(abottom, bbottom);
float c_area = max(cright - cleft, 0.0f) * max(cbottom - ctop, 0.0f);
if (c_area == 0.0f) return 0.0f;
float a_area = max(0.0f, aright - aleft) * max(0.0f, abottom - atop);
float b_area = max(0.0f, bright - bleft) * max(0.0f, bbottom - btop);
return c_area / (a_area + b_area - c_area);
}
__global__ void nms_kernel(float* data, float kNmsThresh, int maxObjects, int numBoxElement){
int position = blockDim.x * blockIdx.x + threadIdx.x;
int count = min((int)data[0], maxObjects);
if (position >= count) return;
// left, top, right, bottom, confidence, class, keepflag
float* pcurrent = data + 1 + position * numBoxElement;
float* pitem;
for (int i = 0; i < count; i++){
pitem = data + 1 + i * numBoxElement;
if (i == position || pcurrent[5] != pitem[5]) continue;
if (pitem[4] >= pcurrent[4]){
if (pitem[4] == pcurrent[4] && i < position) continue;
float iou = box_iou(
pcurrent[0], pcurrent[1], pcurrent[2], pcurrent[3],
pitem[0], pitem[1], pitem[2], pitem[3]
);
if (iou > kNmsThresh){
pcurrent[6] = 0; // 1 = keep, 0 = ignore
return;
}
}
}
}
void nms(float* data, float kNmsThresh, int maxObjects, int numBoxElement, cudaStream_t stream){
int blockSize = maxObjects < 256?maxObjects:256;
int gridSize = (maxObjects + blockSize - 1) / blockSize;
nms_kernel<<<gridSize, blockSize, 0, stream>>>(data, kNmsThresh, maxObjects, numBoxElement);
}
基于c++的推理部署与cuda加速后,相应的模型和代码将作为最终的产出集成进产品里,使用测试集测试确定:
-
单帧推理耗时5.88毫秒,即一秒推理170帧,粗略估计实时处理10路数据问题不大。(还没有发力,指标就达到了,没意思,后面再手动给自己加戏吧,一堆优化策略没上呢)
-
下一步逐步做完性能保持的测试,确保量化部署后的模型和原始模型没有大的精度损失,即可对外发布了。
量化精度损失评估
以最终目标为导向,进行量化后目标精度损失的评估,即这一批数据中每一张图应该出几个框,统计下量化后这些框还在不在,在的话即认为没有精度损失,否则认为有精度损失,简单高效。具体方法如下:
-
逐类别拆分,分析每个类别的量化精度损失;
-
定义IOU阈值为0.5,进行量化前后的目标框匹配,最匹配目标IOU大于0.5即认为匹配上,以量化前作为真值,统计匹配上目标的准确率和召回率,进而判断是否有不可接受的精度损失。
精度保持测试:
-
pytorch模型和c++版本TensorRT模型
有精度损失但差异不大,整体可以用,后续可以着重看损失点在哪里,进行相关优化。
到这为止,YOLOV11的量化部署就告一段路了,达到了基本可用的程度,后续再继续优化,其实整体链路早就打通了,只是在做一致性测试的时候,一直有一些问题,仔细排查后才对齐,而这一块也是做模型部署很重要的一部分工作,后续再着重讲吧。上述工作对应的源码整理中,同样会上传到下面的代码仓库,大家稍作等待。GitHub - luo841297935/vehicle_flowrate
备注:
c++部署部分的代码参考了 GitHub - emptysoal/TensorRT-YOLO11: Based on tensorrt v8.0+, deploy detection, pose, segment, tracking of YOLO11 with C++ and python api.,这里面有较多的yolo系列的工程实现,大家也可以访问支持。
爱罗AI说 是一个和大家交流以计算机视觉为代表的人工智能技术如何在实际生活和工业场景落地使用的公众号,里面包含了我近十年行业公众经验的总结,以及被验证过好用的技术和论文的分享,希望能以一己之力,为技术在行业中真正落地贡献力量,期待和大家的链接。
CSDN上我会放一些纯技术文章,行业经验总结放这边不合适,可以看下面。
历史文章可见公众号
个人微信请添加 Aleo