071、推理引擎 predictor.py 源码逐行解析setup_source到preprocess到inference到postprocess从一次线上推理崩溃说起去年双十一大促前夜我盯着监控面板上突然飙升的推理延迟CPU打满GPU显存却只用了30%。排查到凌晨三点发现是predictor.py里一个不起眼的torch.no_grad()位置放错了——预处理阶段开了梯度计算导致整个pipeline被反向传播的钩子拖死。这种坑只有一行行啃过源码的人才能避开。今天我们就从predictor.py的入口开始把setup_source、preprocess、inference、postprocess这四个阶段拆开揉碎。代码基于YOLOv8的官方实现但核心逻辑在v5、v6、v9里都通用。setup_source别被“source”这个参数骗了defsetup_source(self,source):# 这里踩过坑source可以是str、Path、int、list、np.ndarray# 但千万别传None否则后面imgsz会炸ifisinstance(source,(str,Path)):# 文件或目录路径ifPath(source).is_dir():# 递归扫描所有图片注意过滤隐藏文件self.filessorted(Path(source).rglob(*))self.files[xforxinself.filesifx.suffix.lower()inIMG_FORMATS]else:# 单文件或视频self.files[Path(source)]elifisinstance(source,int):# 摄像头ID这里有个坑0是默认摄像头但有些笔记本摄像头是1self.cam_idsource self.files[]# 别这样写self.files None后面遍历会报TypeErrorelifisinstance(source,list):# 批量图片路径或numpy数组self.filessourceelifisinstance(source,np.ndarray):# 单张图片直接包装成列表self.files[source]else:raiseTypeError(fUnsupported source type:{type(source)})# 初始化数据集加载器这里有个性能关键点self.datasetself._init_dataset(self.files)注意看_init_dataset内部做了什么——它调用了LoadImages或LoadStreams这两个类负责把原始数据转换成模型需要的格式。但有个隐藏逻辑如果source是int摄像头它会启动一个后台线程持续读取帧这个线程如果没处理好会导致主进程卡死在cv2.VideoCapture.read()上。preprocess你以为只是resize太天真了defpreprocess(self,im):# 这里踩过坑im可能是PIL Image、OpenCV BGR、或者已经是tensor# 统一转成RGB的numpy数组ifnotisinstance(im,np.ndarray):imnp.array(im)ifim.shape[2]4:# RGBA转RGBimim[:,:,:3]ifim.shape[2]3andim.dtypenp.uint8:# OpenCV默认BGR但模型训练用的是RGBimim[:,:,::-1]# 别这样写cv2.cvtColor(im, cv2.COLOR_BGR2RGB)慢3倍# 等比例缩放填充保持宽高比h,wim.shape[:2]rmin(self.imgsz/h,self.imgsz/w)# 缩放比例new_h,new_wint(h*r),int(w*r)# 这里有个性能陷阱用cv2.resize比torch.nn.functional.interpolate快imcv2.resize(im,(new_w,new_h),interpolationcv2.INTER_LINEAR)# 填充到正方形注意填充值要和训练时一致dh,dwself.imgsz-new_h,self.imgsz-new_w top,bottomdh//2,dh-dh//2left,rightdw//2,dw-dw//2imcv2.copyMakeBorder(im,top,bottom,left,right,cv2.BORDER_CONSTANT,value(114,114,114))# 归一化转CHW加batch维度imim.astype(np.float32)/255.0imim.transpose(2,0,1)# HWC - CHWimnp.expand_dims(im,axis0)# 加batch维度# 转成tensor注意这里不要用torch.from_numpy().to(device)# 因为后面inference会统一处理deviceimtorch.from_numpy(im)returnim预处理阶段最容易出问题的就是填充逻辑。训练时用的LetterBox和推理时用的copyMakeBorder必须完全一致否则模型会看到奇怪的边缘信息。我见过一个case训练时填充值是0推理时用了114结果小目标检测率直接掉了15%。inferenceforward只是冰山一角definference(self,im,augmentFalse,visualizeFalse):# 这里踩过坑im可能已经在GPU上也可能在CPU上# 统一处理deviceifim.device!self.device:imim.to(self.device)# 开启推理模式关闭梯度计算# 别这样写with torch.no_grad(): 放在外面否则预处理阶段的梯度也会被保留withtorch.no_grad():# 模型前向传播predself.model(im,augmentaugment,visualizevisualize)# 这里有个隐藏逻辑模型可能返回多个输出如v8的detect和segmentifisinstance(pred,(list,tuple)):predpred[0]# 取检测分支# 非极大值抑制注意这里的参数要和训练时匹配prednon_max_suppression(pred,conf_thresself.conf_thres,iou_thresself.iou_thres,classesself.classes,agnosticself.agnostic_nms,max_detself.max_det,)returnpredNMS这一步是推理的瓶颈。YOLOv8默认用的是torchvision.ops.nms但如果你用的是v5或v6它们用的是自定义的C实现。这里有个性能优化点如果batch size大于1可以先把所有预测结果拼在一起做一次NMS再按batch拆分能减少重复计算。postprocess把坐标映射回原图defpostprocess(self,preds,im,orig_imgs):# 这里踩过坑preds是list of tensors每个tensor对应一张图# 但orig_imgs可能是list of arrays也可能是单张arrayresults[]fori,predinenumerate(preds):iflen(pred)0:results.append([])continue# 获取原始图像尺寸ifisinstance(orig_imgs,list):orig_h,orig_worig_imgs[i].shape[:2]else:orig_h,orig_worig_imgs.shape[:2]# 获取预处理后的尺寸填充后的h,wim.shape[2:]# 注意im是CHW格式# 计算缩放和填充的逆变换rmin(h/orig_h,w/orig_w)new_h,new_wint(orig_h*r),int(orig_w*r)dh(h-new_h)/2dw(w-new_w)/2# 将预测框从模型坐标系映射回原图坐标系# 别这样写直接除以缩放比例会忽略填充偏移pred[:,[0,2]](pred[:,[0,2]]-dw)/r# x坐标pred[:,[1,3]](pred[:,[1,3]]-dh)/r# y坐标# 裁剪到原图边界pred[:,[0,2]]pred[:,[0,2]].clamp(0,orig_w)pred[:,[1,3]]pred[:,[1,3]].clamp(0,orig_h)# 转成整数坐标可选predpred.round().int()results.append(pred)returnresults后处理最容易翻车的地方是坐标映射。很多人直接pred / r忽略了填充的偏移量。更隐蔽的问题是如果原图是矩形填充后模型预测的框可能有一部分在填充区域里裁剪到原图边界后框的面积会变小导致置信度虚高。个人经验性建议预处理和后处理的参数必须和训练时完全一致。我见过最离谱的bug是训练时用了interpolationcv2.INTER_CUBIC推理时用了INTER_LINEAR结果模型对边缘纹理的响应完全不同。建议把预处理参数写成一个config类训练和推理共用。NMS的conf_thres和iou_thres不要用固定值。线上环境光照变化大建议根据场景动态调整。比如夜间场景降低conf_thres白天场景提高。我写了个自适应阈值模块根据检测框的平均置信度自动调节效果比固定值好10%。batch推理时注意内存碎片。如果每张图尺寸不同预处理后的tensor大小也不同频繁分配释放会导致内存碎片。建议统一resize到固定尺寸或者用torch.cuda.empty_cache()定期清理。别在GPU上做后处理。坐标映射、裁剪这些操作在CPU上更快而且不会占用显存。我见过有人把整个postprocess放在GPU上结果显存爆了推理速度反而慢了。监控推理pipeline的每个阶段耗时。用torch.cuda.Event记录时间或者用time.perf_counter()。我通常会在setup_source、preprocess、inference、postprocess四个阶段各打一个时间戳这样能快速定位瓶颈。有一次发现preprocess占了60%的时间排查后发现是cv2.resize的interpolation参数设成了INTER_CUBIC改成INTER_LINEAR后直接快了3倍。最后说一句predictor.py是整个推理系统的核心每一行代码都值得反复推敲。别想着“先跑通再说”线上出问题的时候你连睡觉的时间都没有。