1. 为什么Pytorch部署是Yolov5落地的第一道坎搞过深度学习项目的人都有一个共识:训练出模型只是万里长征第一步,真正把模型塞进实际业务里跑起来,才是噩梦的开始。Yolov5作为目标检测领域最流行的框架之一,训练脚本写得再溜,到了部署环节照样有一堆人卡住。这个系列教程的第二篇,专门聊Yolov5的Pytorch部署,原因很简单——Pytorch部署是所有部署方案里最“原汁原味”的一条路,它不需要格式转换,不需要额外推理引擎,直接拿训练好的权重就能跑,适合快速验证、小规模服务、以及作为其他部署方案(ONNX、TensorRT、OpenVINO)的基准对照。但“简单”不等于“没坑”。我见过太多人拿着detect.py跑通了就以为部署完成了,结果一上生产环境,batch推理不会写、GPU显存爆了、预处理和后处理对不上、多线程下模型输出乱序……这些问题在demo里永远不会暴露,但到了真实场景里一个比一个致命。这篇文章要做的,就是把Yolov5的Pytorch部署从“能跑”推到“能用”的层面,把那些官方文档里一笔带过、但实际项目中必须处理的细节全部摊开讲。适合谁看?如果你已经跑通过Yolov5的训练,手里有.pt权重文件,现在需要把它集成到一个Python服务里,或者需要写一个独立的推理脚本嵌入到现有系统中,那这篇内容就是为你准备的。如果你还没训练过模型,建议先去看本系列第一篇,把训练流程走通再回来。整篇内容会围绕一个核心目标展开:让你拿到任意一个Yolov5的.pt权重,都能写出一个结构清晰、性能可控、易于集成的Pytorch推理模块。2. 部署前的环境梳理与权重准备2.1 环境依赖的精确锁定Pytorch部署Yolov5的第一个坑,往往不是代码问题,而是环境问题。Yolov5的代码仓库更新频率很高,不同版本对Pytorch、TorchVision、Python的版本要求不一样,而Pytorch本身又和CUDA驱动版本强绑定。我踩过最典型的一次坑是:训练环境用的是Pytorch 1.10 + CUDA 11.3,部署服务器上装的是Pytorch 2.0 + CUDA 11.8,结果加载权重时直接报了一堆KeyError,原因是模型定义里的某些层命名在新版本里变了。所以第一步不是急着写代码,而是把环境锁死。我的建议是:部署环境的Pytorch版本尽量和训练环境保持一致,至少保证大版本号相同。如果实在无法一致,那就要做好权重加载时做键名映射的准备。具体操作上,先确认训练时用的版本:pip show torch torchvision输出里会显示具体的版本号,比如torch 1.13.1+cu117,这个cu117就代表编译时用的CUDA版本是11.7。部署环境最好装同样的版本。如果服务器CUDA驱动版本较低,装不了高版本Pytorch,那就只能降级Pytorch,但要注意Yolov5某些新特性(比如某些注意力模块)可能在低版本Pytorch上不支持。另外,Yolov5的依赖不止Pytorch。requirements.txt里还有numpy、opencv-python、pillow、pyyaml、requests、scipy、tqdm、matplotlib、pandas、seaborn这些。部署环境里其实只需要其中一部分:numpy用于数组运算,opencv-python用于图像读写和预处理,pillow用于某些图像格式兼容,pyyaml用于读取模型配置,tqdm用于进度条(可选)。matplotlib、pandas、seaborn这些是训练和可视化用的,部署时完全可以不装,能省不少空间。注意:opencv-python和opencv-python-headless的区别。如果你的部署环境没有图形界面(比如Docker容器或远程服务器),一定要装headless版本,否则会报libGL.so.1找不到的错误。这个坑我见过至少十个人踩过。2.2 权重文件的获取与校验Yolov5的权重文件通常有三种来源:官方预训练权重(yolov5s.pt、yolov5m.pt等)、自己训练保存的best.pt、以及从last.pt恢复的权重。不管哪种,拿到权重后第一件事是校验它能不能正常加载。很多人直接拿last.pt去部署,结果发现模型还在训练中途的状态,性能远不如best.pt。所以部署前务必确认用的是best.pt。校验的方法很简单,用Pytorch加载一下:import torch ckpt = torch.load('best.pt', map_location='cpu') print(type(ckpt)) print(ckpt.keys() if isinstance(ckpt, dict) else 'not a dict')Yolov5保存的权重通常是一个字典,里面包含model(模型对象)、ema(指数移动平均)、optimizer(优化器状态)、best_fitness、epoch等信息。部署时我们只关心model和ema。如果ema存在,优先用ema,因为它的性能通常更稳定。加载到CPU上是为了避免在没有GPU的机器上直接报错,后续再根据实际情况转移到GPU。还有一个细节:Yolov5的权重里保存的model是一个DetectionModel对象,它包含了模型结构定义。这意味着你不需要额外提供yaml配置文件就能重建模型。但这也带来一个问题:如果部署环境的Yolov5代码版本和训练时不一致,DetectionModel类的定义可能不同,导致加载失败。所以最稳妥的做法是:把训练时用的Yolov5代码仓库一起打包到部署环境,或者至少保证models目录下的代码一致。2.3 模型结构的快速验证加载完权重后,别急着跑推理,先做一次前向传播验证。构造一个随机输入,看看输出形状对不对:import torch model = ckpt['ema' if ckpt.get('ema') else 'model'].float() model.eval() dummy = torch.randn(1, 3, 640, 640) with torch.no_grad(): out = model(dummy) print(out[0].shape)Yolov5的输出是一个列表,推理时通常只取第一个元素,形状是[batch, num_anchors, 5 + num_classes]。以COCO数据集80类为例,num_anchors在640输入下通常是25200,所以输出形状是[1, 25200, 85]。这个85就是4个坐标 + 1个置信度 + 80个类别概率。如果你看到输出形状不对,比如少了维度或者多了维度,那说明模型加载有问题,或者输入尺寸不对。这一步的意义在于:在写复杂推理逻辑之前,先确认模型本身是好的。我遇到过有人把分类模型的权重当成检测模型加载,结果输出形状完全对不上,排查了半天才发现拿错了文件。3. 推理流程的完整拆解与代码实现3.1 预处理:从原始图像到模型输入Yolov5的预处理看起来简单,但细节很多。官方detect.py里的预处理逻辑封装在LoadImages和letterbox函数里,核心步骤是:保持长宽比的缩放 + 填充到32的倍数。为什么是32?因为Yolov5的Backbone有5次下采样,每次下采样2倍,总共32倍。如果输入尺寸不是32的倍数,特征图尺寸会对不齐,导致拼接失败。letterbox的具体做法是:计算缩放比例r = min(target_h / h, target_w / w),然后用cv2.resize按这个比例缩放,最后把缩放后的图像粘贴到一个灰色(114, 114, 114)的画布上,画布尺寸是target_h x target_w。填充的位置是居中的,左右或上下均匀分配。这里有个容易忽略的点:填充值为什么是114?这是Yolov5作者选的,来源于ImageNet的均值。实际上用0填充也能跑,但114是训练时用的值,部署时保持一致能避免分布偏移。我实测过,用0填充和114填充,在大多数场景下差异很小,但在一些边缘场景(比如图像边缘有目标)可能会有几个百分点的mAP差异。预处理还有一个关键点:BGR到RGB的转换。OpenCV读进来的是BGR,而Pytorch模型训练时用的是RGB。所以cv2.imread之后要cv2.cvtColor(img, cv2.COLOR_BGR2RGB)。这个顺序不能错,否则颜色通道对不上,检测结果会莫名其妙地差。归一化方面,Yolov5的预处理是把像素值从0-255缩放到0-1,即img / 255.0。然后转成CHW格式,再增加batch维度。整个预处理的代码可以写成这样:import cv2 import numpy as np import torch def letterbox(img, new_shape=640, color=(114, 114, 114)): shape = img.shape[:2] r = min(new_shape / shape[0], new_shape / shape[1]) new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) dw, dh = new_shape - new_unpad[0], new_shape - new_unpad[1] dw, dh = dw // 2, dh // 2 img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = dh, new_shape - new_unpad[1] - dh left, right = dw, new_shape - new_unpad[0] - dw img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) return img, r, (dw, dh) def preprocess(img_path, img_size=640): img0 = cv2.imread(img_path) img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB) img, ratio, pad = letterbox(img0, img_size) img = img.transpose(2, 0, 1) # HWC -> CHW img = np.ascontiguousarray(img) img = torch.from_numpy(img).float() / 255.0 img = img.unsqueeze(0) return img, img0, ratio, pad这段代码里,ratio和pad要保留下来,后处理时需要用它们把检测框映射回原图坐标。很多人后处理时忘了这一步,结果画出来的框位置全错。3.2 推理:模型前向传播与GPU加速推理这一步本身很简单,就是model(img)。但有几个性能相关的点值得展开。第一,torch.no_grad()必须加。不加的话,Pytorch会构建计算图,显存占用会暴涨,推理速度也会慢很多。这个坑很基础,但确实有人忘。第二,半精度推理。如果用的是支持FP16的GPU(比如图灵架构及以后的N卡),可以把模型和输入都转成half(),推理速度能提升30%-50%,显存占用减半。但要注意,某些操作在FP16下可能溢出,导致输出NaN。我的做法是:先用FP32跑一遍,确认结果正常,再切FP16对比,如果结果差异在可接受范围内(比如框的坐标差异小于1像素),就用FP16。model = model.half().cuda() img = img.half().cuda() with torch.no_grad(): pred = model(img)[0]第三,batch推理。如果有多张图要处理,尽量拼成一个batch。Yolov5的模型支持动态batch,但前提是所有输入尺寸一致。如果图片尺寸不一,要么统一resize到同一尺寸(会损失精度),要么逐张推理。实际项目中,我通常会把输入尺寸固定(比如640x640),然后按batch=8或16来跑,吞吐量比逐张推理高好几倍。第四,model.eval()不能忘。虽然Yolov5的推理模式下BN层和Dropout的行为和训练模式不同,但如果不调eval(),BN层会用batch统计量而不是滑动平均,结果会不稳定。这个坑在单张推理时可能不明显,但batch推理时差异会放大。3.3 后处理:从25200个预测框到最终结果后处理是Yolov5部署里最复杂的部分。模型输出[1, 25200, 85],这25200个框里绝大多数是背景或低置信度的,需要经过置信度过滤、NMS、坐标映射三步才能得到最终结果。置信度过滤很简单:conf = obj_conf * cls_conf,Yolov5的85维里,第4维是objectness,第5维开始是类别概率。最终置信度是两者相乘。然后设定一个阈值(通常0.25),过滤掉低置信度的框。NMS(非极大值抑制)是核心。Yolov5官方用的是torchvision.ops.nms,输入是框的坐标和分数,输出是保留的索引。这里有个细节:NMS的IoU阈值。默认是0.45,但实际项目中要根据场景调整。比如密集场景(人群检测)要调低到0.3-0.4,避免漏检;稀疏场景可以调到0.5-0.6,避免误删。坐标映射是最后一步。模型输出的框坐标是相对于letterbox后的图像的,需要先减去padding,再除以ratio,才能映射回原图。具体公式是:boxes[:, [0, 2]] -= pad[0] boxes[:, [1, 3]] -= pad[1] boxes[:, :4] /= ratio然后还要做clip,把坐标限制在原图范围内。这一步不做的话,画框时可能画出图像边界。完整的后处理代码:def non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45): pred = pred[0] # [25200, 85] conf = pred[:, 4] * pred[:, 5:].max(1)[0] mask = conf > conf_thres pred = pred[mask] if pred.shape[0] == 0: return None box = xywh2xyxy(pred[:, :4]) scores = pred[:, 4] * pred[:, 5:].max(1)[0] cls = pred[:, 5:].argmax(1) keep = torchvision.ops.nms(box, scores, iou_thres) return torch.cat([box[keep], scores[keep].unsqueeze(1), cls[keep].unsqueeze(1).float()], 1)xywh2xyxy是把中心点+宽高的格式转成左上角+右下角的格式,这是NMS需要的输入格式。这个转换看起来简单,但写错的人不少,尤其是x1 = x - w/2和x2 = x + w/2这两个公式,容易把w和h搞混。3.4 完整推理脚本的组装把预处理、推理、后处理串起来,就是一个完整的推理脚本。我习惯把它封装成一个类,方便复用:class Yolov5Detector: def __init__(self, weights, device='cuda', img_size=640, conf_thres=0.25, iou_thres=0.45): self.device = device self.img_size = img_size self.conf_thres = conf_thres self.iou_thres = iou_thres ckpt = torch.load(weights, map_location='cpu') self.model = ckpt['ema' if ckpt.get('ema') else 'model'].float() self.model.eval() self.model.to(device) self.names = ckpt['model'].names if hasattr(ckpt['model'], 'names') else None def detect(self, img_path): img, img0, ratio, pad = preprocess(img_path, self.img_size) img = img.to(self.device) with torch.no_grad(): pred = self.model(img)[0] det = non_max_suppression(pred, self.conf_thres, self.iou_thres) if det is None: return [] det[:, [0, 2]] -= pad[0] det[:, [1, 3]] -= pad[1] det[:, :4] /= ratio det[:, [0, 2]] = det[:, [0, 2]].clamp(0, img0.shape[1]) det[:, [1, 3]] = det[:, [1, 3]].clamp(0, img0.shape[0]) results = [] for *box, conf, cls in det: results.append({ 'box': [int(x) for x in box], 'conf': float(conf), 'cls': int(cls), 'name': self.names[int(cls)] if self.names else str(int(cls)) }) return results这个类可以直接拿去用,输入图片路径,输出检测结果列表。实际项目中,我会把img_path换成img数组,避免重复读盘。4. 性能优化与生产环境适配4.1 显存管理与批处理策略Pytorch部署Yolov5时,显存是稀缺资源。一张1080Ti有11GB显存,跑Yolov5s在640输入下单张推理大概占1.5GB,但如果你开batch=16,显存占用会飙升到8GB以上。所以batch size不能随便设,要根据显存和模型大小来算。一个粗略的估算公式:显存占用 ≈ 模型参数显存 + 激活值显存 × batch_size。Yolov5s的参数约7M,FP32下占28MB,FP16下占14MB,这部分很小。大头是激活值,也就是前向传播过程中每一层的输出。Yolov5s在640输入下,单张图的激活值约1.2GB(FP32),所以batch=8时约9.6GB,加上模型和其他开销,11GB显存基本就满了。实际部署时,我通常会用torch.cuda.memory_allocated()来监控显存,动态调整batch size。如果显存不够,优先考虑:1)切FP16,显存直接减半;2)降低输入尺寸,比如从640降到416,激活值按平方比例减少;3)用梯度累积模拟大batch(但推理时不需要梯度,所以这条不适用)。还有一个技巧:推理时用torch.inference_mode()替代torch.no_grad()。前者是Pytorch 1.9引入的,比后者更彻底,能进一步减少显存占用和提升速度。实测下来,差异在5%左右,但积少成多。4.2 多线程与异步推理生产环境里,推理服务通常要同时处理多个请求。如果每个请求都串行跑,吞吐量上不去。Pytorch本身不是线程安全的,但推理时如果模型是只读的,多线程调用model()是可以的。不过GIL的存在会让Python层面的多线程效果打折扣,所以更推荐用多进程或者异步IO + 单线程推理。我常用的方案是:用一个队列接收请求,一个后台线程专门跑推理,推理结果通过另一个队列返回。这样推理线程可以连续跑,不用等IO。如果请求量很大,就起多个进程,每个进程加载一份模型,用共享内存或消息队列分发请求。这个方案在4核CPU + 1张GPU的机器上,能跑到每秒50-80张图(Yolov5s,640输入,FP16)。注意:多进程加载模型时,每个进程都会占一份显存。如果显存不够,可以用torch.multiprocessing的spawn模式,配合torch.cuda.set_device()让每个进程用不同的GPU。但单GPU下,多进程反而会因为显存竞争变慢,这时候多线程 + 队列更合适。4.3 模型量化与加速的取舍Pytorch支持动态量化和静态量化,能把FP32模型转成INT8,显存占用减少75%,速度提升2-3倍。但Yolov5的量化不是开箱即用的,因为检测头里的某些操作(比如sigmoid、exp)对量化敏感,直接量化会导致精度大幅下降。我试过用Pytorch的torch.quantization.quantize_dynamic对Yolov5做动态量化,结果mAP掉了15个百分点,基本不可用。后来改用量化感知训练(QAT),在训练时插入伪量化节点,让模型适应量化误差,最终mAP只掉了2个百分点,速度提升了2.5倍。但QAT需要重新训练,成本较高,适合对速度要求极高的场景。如果不想重新训练,又想加速,可以考虑TorchScript。把模型转成TorchScript后,Pytorch会做图优化,去掉一些冗余操作,速度能提升10%-20%。转换方法很简单:scripted_model = torch.jit.script(model) scripted_model.save('yolov5s_scripted.pt')但TorchScript对动态控制流支持有限,Yolov5的某些版本可能转不过去。我实测Yolov5 v6.0可以转,v7.0有些版本会报错。转换前先用torch.jit.trace试一下,如果trace能过,script通常也能过。4.4 服务化部署的接口设计如果要把Yolov5的Pytorch推理做成一个服务,接口设计要考虑几点:输入格式、输出格式、错误处理、并发控制。输入格式我通常用JSON,包含image字段(base64编码的图片)和可选参数(conf_thres、iou_thres)。输出也是JSON,包含boxes、scores、classes三个数组。这样前端和后端都好处理。错误处理要覆盖:图片解码失败、模型推理异常、显存不足。显存不足时,可以捕获RuntimeError,然后清空torch.cuda.empty_cache(),再重试一次。如果还不行,就返回503,让客户端稍后重试。并发控制用信号量或者线程池。我一般用concurrent.futures.ThreadPoolExecutor,设置max_workers等于GPU能承受的并发数。每个worker里跑一次推理,结果通过Future返回。这样既能控制并发,又能复用线程。5. 常见问题排查与实战避坑指南5.1 加载权重报错:KeyError与Unexpected key这是最常见的问题,没有之一。报错信息通常是Missing key(s) in state_dict或者Unexpected key(s) in state_dict。原因一般是训练和部署的Yolov5代码版本不一致,导致模型定义里的层名变了。解决方法有两种:一是统一代码版本,把训练时的models目录复制到部署环境;二是手动映射键名,用collections.OrderedDict重新组织state_dict。第二种方法麻烦但灵活,适合无法统一版本的场景。具体操作:先打印训练时权重的键名和部署时模型的键名,对比差异,然后写一个映射函数。比如训练时叫model.0.conv.weight,部署时叫model.0.conv.0.weight,那就把前者映射到后者。这个工作很枯燥,但一次写好就能复用。5.2 推理结果为空或框位置错乱推理结果为空,通常是置信度阈值设太高,或者预处理时颜色通道搞错了。先检查conf_thres是不是0.25以上,再检查cv2.cvtColor有没有加。如果框位置错乱,比如框整体偏移或者大小不对,那多半是后处理的坐标映射出了问题。检查pad和ratio的计算是否正确,尤其是pad的左右和上下顺序。还有一个隐蔽的坑:letterbox的padding顺序。cv2.copyMakeBorder的参数顺序是top, bottom, left, right,而计算dw和dh时,dw是宽度的padding,对应left和right;dh是高度的padding,对应top和bottom。如果搞反了,框会整体偏移。我见过有人把dw和dh搞反,结果框全部偏到右下角。5.3 显存泄漏与OOM排查显存泄漏的表现是:跑一段时间后,显存占用越来越高,最终OOM。原因通常是:1)没有用torch.no_grad(),计算图没释放;2)中间变量(比如pred、det)一直持有引用,没被GC回收;3)多线程下,每个线程都创建了新的CUDA上下文。排查方法:用torch.cuda.memory_summary()看显存分配情况,找出哪个部分在增长。如果是计算图没释放,加torch.no_grad();如果是变量引用,手动del并调torch.cuda.empty_cache();如果是多线程上下文,改用单线程推理或者torch.multiprocessing。提示:torch.cuda.empty_cache()只能释放未被引用的显存,如果变量还被引用,调了也没用。所以关键是找到并删除引用。5.4 常见问题速查表问题现象可能原因排查方法解决方案加载权重报KeyError代码版本不一致对比键名统一版本或手动映射推理结果为空置信度阈值过高打印原始输出降低conf_thres框位置偏移padding顺序错误检查dw/dh修正copyMakeBorder参数显存持续增长计算图未释放memory_summary加no_grad,删引用推理速度慢未用FP16检查dtype转half,用inference_mode多线程结果乱序共享变量竞争加锁或改队列用队列串行化颜色通道错误BGR未转RGB检查cvtColor加BGR2RGB转换输入尺寸不对齐非32倍数检查img_size调整为32的倍数5.5 独家避坑经验第一个经验:永远保留一份“黄金输入”。就是一张固定的测试图,每次改完代码都跑一遍,看输出是否一致。这样能快速发现回归问题。我习惯把这张图的检测结果(框坐标、分数)存成JSON,改代码后对比差异,超过阈值就报警。第二个经验:预处理和后处理要成对测试。单独测预处理,看输出的tensor形状和数值范围;单独测后处理,用固定的tensor输入,看输出的框是否正确。这样能把问题定位到具体环节,而不是笼统地“推理不对”。第三个经验:GPU和CPU的结果要对比。有时候GPU上跑出来正常,CPU上跑出来不对,原因是某些操作在CPU和GPU上的实现有细微差异(比如插值算法)。如果部署环境可能没有GPU,一定要在CPU上验证一遍。第四个经验:模型切换要热更新。生产环境里,模型可能会更新。如果每次更新都重启服务,用户体验很差。我的做法是:把模型加载封装成一个函数,用文件监听或者API触发重新加载。加载新模型时,先加载到临时变量,验证通过后再替换旧模型,避免加载失败导致服务不可用。6. 从Pytorch部署到其他方案的过渡Pytorch部署跑通之后,你可能会遇到性能瓶颈,比如单张推理要50ms,但业务要求10ms。这时候就需要考虑其他部署方案了。Pytorch部署的最大价值,其实是作为基准:它告诉你模型在理想情况下的精度上限,后续任何加速方案(ONNX、TensorRT、OpenVINO)都要和它对比,确保精度损失在可接受范围内。从Pytorch到ONNX的转换,核心是torch.onnx.export。但Yolov5的导出有几个坑:1)输出需要是tuple而不是list;2)动态batch需要设置dynamic_axes;3)某些算子(比如torch.nn.functional.interpolate)在ONNX里的实现可能和Pytorch不同。我通常会在导出后,用onnxruntime跑一遍,和Pytorch的结果对比,确保差异小于1e-3。从ONNX到TensorRT,又是另一个坑。TensorRT对算子支持有限,Yolov5的某些版本需要插件(plugin)才能转。而且TensorRT的精度和Pytorch可能有差异,尤其是FP16模式下。我的建议是:先用Pytorch部署验证业务逻辑,再逐步替换推理后端。这样即使加速方案出问题,也有一个可用的兜底方案。最后分享一个我个人的习惯:把Pytorch推理脚本和模型权重一起打包成Docker镜像。这样部署时不用操心环境问题,直接docker run就能跑。镜像里只装必要的依赖(Pytorch、OpenCV、NumPy),体积能控制在2GB以内。如果显存紧张,还可以用nvidia-docker限制GPU内存,避免影响其他服务。
1. 为什么Pytorch部署是Yolov5落地的第一道坎搞过深度学习项目的人都有一个共识:训练出模型只是万里长征第一步,真正把模型塞进实际业务里跑起来,才是噩梦的开始。Yolov5作为目标检测领域最流行的框架之一,训练脚本写得再溜,到了部署环节照样有一堆人卡住。这个系列教程的第二篇,专门聊Yolov5的Pytorch部署,原因很简单——Pytorch部署是所有部署方案里最“原汁原味”的一条路,它不需要格式转换,不需要额外推理引擎,直接拿训练好的权重就能跑,适合快速验证、小规模服务、以及作为其他部署方案(ONNX、TensorRT、OpenVINO)的基准对照。但“简单”不等于“没坑”。我见过太多人拿着detect.py跑通了就以为部署完成了,结果一上生产环境,batch推理不会写、GPU显存爆了、预处理和后处理对不上、多线程下模型输出乱序……这些问题在demo里永远不会暴露,但到了真实场景里一个比一个致命。这篇文章要做的,就是把Yolov5的Pytorch部署从“能跑”推到“能用”的层面,把那些官方文档里一笔带过、但实际项目中必须处理的细节全部摊开讲。适合谁看?如果你已经跑通过Yolov5的训练,手里有.pt权重文件,现在需要把它集成到一个Python服务里,或者需要写一个独立的推理脚本嵌入到现有系统中,那这篇内容就是为你准备的。如果你还没训练过模型,建议先去看本系列第一篇,把训练流程走通再回来。整篇内容会围绕一个核心目标展开:让你拿到任意一个Yolov5的.pt权重,都能写出一个结构清晰、性能可控、易于集成的Pytorch推理模块。2. 部署前的环境梳理与权重准备2.1 环境依赖的精确锁定Pytorch部署Yolov5的第一个坑,往往不是代码问题,而是环境问题。Yolov5的代码仓库更新频率很高,不同版本对Pytorch、TorchVision、Python的版本要求不一样,而Pytorch本身又和CUDA驱动版本强绑定。我踩过最典型的一次坑是:训练环境用的是Pytorch 1.10 + CUDA 11.3,部署服务器上装的是Pytorch 2.0 + CUDA 11.8,结果加载权重时直接报了一堆KeyError,原因是模型定义里的某些层命名在新版本里变了。所以第一步不是急着写代码,而是把环境锁死。我的建议是:部署环境的Pytorch版本尽量和训练环境保持一致,至少保证大版本号相同。如果实在无法一致,那就要做好权重加载时做键名映射的准备。具体操作上,先确认训练时用的版本:pip show torch torchvision输出里会显示具体的版本号,比如torch 1.13.1+cu117,这个cu117就代表编译时用的CUDA版本是11.7。部署环境最好装同样的版本。如果服务器CUDA驱动版本较低,装不了高版本Pytorch,那就只能降级Pytorch,但要注意Yolov5某些新特性(比如某些注意力模块)可能在低版本Pytorch上不支持。另外,Yolov5的依赖不止Pytorch。requirements.txt里还有numpy、opencv-python、pillow、pyyaml、requests、scipy、tqdm、matplotlib、pandas、seaborn这些。部署环境里其实只需要其中一部分:numpy用于数组运算,opencv-python用于图像读写和预处理,pillow用于某些图像格式兼容,pyyaml用于读取模型配置,tqdm用于进度条(可选)。matplotlib、pandas、seaborn这些是训练和可视化用的,部署时完全可以不装,能省不少空间。注意:opencv-python和opencv-python-headless的区别。如果你的部署环境没有图形界面(比如Docker容器或远程服务器),一定要装headless版本,否则会报libGL.so.1找不到的错误。这个坑我见过至少十个人踩过。2.2 权重文件的获取与校验Yolov5的权重文件通常有三种来源:官方预训练权重(yolov5s.pt、yolov5m.pt等)、自己训练保存的best.pt、以及从last.pt恢复的权重。不管哪种,拿到权重后第一件事是校验它能不能正常加载。很多人直接拿last.pt去部署,结果发现模型还在训练中途的状态,性能远不如best.pt。所以部署前务必确认用的是best.pt。校验的方法很简单,用Pytorch加载一下:import torch ckpt = torch.load('best.pt', map_location='cpu') print(type(ckpt)) print(ckpt.keys() if isinstance(ckpt, dict) else 'not a dict')Yolov5保存的权重通常是一个字典,里面包含model(模型对象)、ema(指数移动平均)、optimizer(优化器状态)、best_fitness、epoch等信息。部署时我们只关心model和ema。如果ema存在,优先用ema,因为它的性能通常更稳定。加载到CPU上是为了避免在没有GPU的机器上直接报错,后续再根据实际情况转移到GPU。还有一个细节:Yolov5的权重里保存的model是一个DetectionModel对象,它包含了模型结构定义。这意味着你不需要额外提供yaml配置文件就能重建模型。但这也带来一个问题:如果部署环境的Yolov5代码版本和训练时不一致,DetectionModel类的定义可能不同,导致加载失败。所以最稳妥的做法是:把训练时用的Yolov5代码仓库一起打包到部署环境,或者至少保证models目录下的代码一致。2.3 模型结构的快速验证加载完权重后,别急着跑推理,先做一次前向传播验证。构造一个随机输入,看看输出形状对不对:import torch model = ckpt['ema' if ckpt.get('ema') else 'model'].float() model.eval() dummy = torch.randn(1, 3, 640, 640) with torch.no_grad(): out = model(dummy) print(out[0].shape)Yolov5的输出是一个列表,推理时通常只取第一个元素,形状是[batch, num_anchors, 5 + num_classes]。以COCO数据集80类为例,num_anchors在640输入下通常是25200,所以输出形状是[1, 25200, 85]。这个85就是4个坐标 + 1个置信度 + 80个类别概率。如果你看到输出形状不对,比如少了维度或者多了维度,那说明模型加载有问题,或者输入尺寸不对。这一步的意义在于:在写复杂推理逻辑之前,先确认模型本身是好的。我遇到过有人把分类模型的权重当成检测模型加载,结果输出形状完全对不上,排查了半天才发现拿错了文件。3. 推理流程的完整拆解与代码实现3.1 预处理:从原始图像到模型输入Yolov5的预处理看起来简单,但细节很多。官方detect.py里的预处理逻辑封装在LoadImages和letterbox函数里,核心步骤是:保持长宽比的缩放 + 填充到32的倍数。为什么是32?因为Yolov5的Backbone有5次下采样,每次下采样2倍,总共32倍。如果输入尺寸不是32的倍数,特征图尺寸会对不齐,导致拼接失败。letterbox的具体做法是:计算缩放比例r = min(target_h / h, target_w / w),然后用cv2.resize按这个比例缩放,最后把缩放后的图像粘贴到一个灰色(114, 114, 114)的画布上,画布尺寸是target_h x target_w。填充的位置是居中的,左右或上下均匀分配。这里有个容易忽略的点:填充值为什么是114?这是Yolov5作者选的,来源于ImageNet的均值。实际上用0填充也能跑,但114是训练时用的值,部署时保持一致能避免分布偏移。我实测过,用0填充和114填充,在大多数场景下差异很小,但在一些边缘场景(比如图像边缘有目标)可能会有几个百分点的mAP差异。预处理还有一个关键点:BGR到RGB的转换。OpenCV读进来的是BGR,而Pytorch模型训练时用的是RGB。所以cv2.imread之后要cv2.cvtColor(img, cv2.COLOR_BGR2RGB)。这个顺序不能错,否则颜色通道对不上,检测结果会莫名其妙地差。归一化方面,Yolov5的预处理是把像素值从0-255缩放到0-1,即img / 255.0。然后转成CHW格式,再增加batch维度。整个预处理的代码可以写成这样:import cv2 import numpy as np import torch def letterbox(img, new_shape=640, color=(114, 114, 114)): shape = img.shape[:2] r = min(new_shape / shape[0], new_shape / shape[1]) new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r))) dw, dh = new_shape - new_unpad[0], new_shape - new_unpad[1] dw, dh = dw // 2, dh // 2 img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = dh, new_shape - new_unpad[1] - dh left, right = dw, new_shape - new_unpad[0] - dw img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) return img, r, (dw, dh) def preprocess(img_path, img_size=640): img0 = cv2.imread(img_path) img0 = cv2.cvtColor(img0, cv2.COLOR_BGR2RGB) img, ratio, pad = letterbox(img0, img_size) img = img.transpose(2, 0, 1) # HWC -> CHW img = np.ascontiguousarray(img) img = torch.from_numpy(img).float() / 255.0 img = img.unsqueeze(0) return img, img0, ratio, pad这段代码里,ratio和pad要保留下来,后处理时需要用它们把检测框映射回原图坐标。很多人后处理时忘了这一步,结果画出来的框位置全错。3.2 推理:模型前向传播与GPU加速推理这一步本身很简单,就是model(img)。但有几个性能相关的点值得展开。第一,torch.no_grad()必须加。不加的话,Pytorch会构建计算图,显存占用会暴涨,推理速度也会慢很多。这个坑很基础,但确实有人忘。第二,半精度推理。如果用的是支持FP16的GPU(比如图灵架构及以后的N卡),可以把模型和输入都转成half(),推理速度能提升30%-50%,显存占用减半。但要注意,某些操作在FP16下可能溢出,导致输出NaN。我的做法是:先用FP32跑一遍,确认结果正常,再切FP16对比,如果结果差异在可接受范围内(比如框的坐标差异小于1像素),就用FP16。model = model.half().cuda() img = img.half().cuda() with torch.no_grad(): pred = model(img)[0]第三,batch推理。如果有多张图要处理,尽量拼成一个batch。Yolov5的模型支持动态batch,但前提是所有输入尺寸一致。如果图片尺寸不一,要么统一resize到同一尺寸(会损失精度),要么逐张推理。实际项目中,我通常会把输入尺寸固定(比如640x640),然后按batch=8或16来跑,吞吐量比逐张推理高好几倍。第四,model.eval()不能忘。虽然Yolov5的推理模式下BN层和Dropout的行为和训练模式不同,但如果不调eval(),BN层会用batch统计量而不是滑动平均,结果会不稳定。这个坑在单张推理时可能不明显,但batch推理时差异会放大。3.3 后处理:从25200个预测框到最终结果后处理是Yolov5部署里最复杂的部分。模型输出[1, 25200, 85],这25200个框里绝大多数是背景或低置信度的,需要经过置信度过滤、NMS、坐标映射三步才能得到最终结果。置信度过滤很简单:conf = obj_conf * cls_conf,Yolov5的85维里,第4维是objectness,第5维开始是类别概率。最终置信度是两者相乘。然后设定一个阈值(通常0.25),过滤掉低置信度的框。NMS(非极大值抑制)是核心。Yolov5官方用的是torchvision.ops.nms,输入是框的坐标和分数,输出是保留的索引。这里有个细节:NMS的IoU阈值。默认是0.45,但实际项目中要根据场景调整。比如密集场景(人群检测)要调低到0.3-0.4,避免漏检;稀疏场景可以调到0.5-0.6,避免误删。坐标映射是最后一步。模型输出的框坐标是相对于letterbox后的图像的,需要先减去padding,再除以ratio,才能映射回原图。具体公式是:boxes[:, [0, 2]] -= pad[0] boxes[:, [1, 3]] -= pad[1] boxes[:, :4] /= ratio然后还要做clip,把坐标限制在原图范围内。这一步不做的话,画框时可能画出图像边界。完整的后处理代码:def non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45): pred = pred[0] # [25200, 85] conf = pred[:, 4] * pred[:, 5:].max(1)[0] mask = conf > conf_thres pred = pred[mask] if pred.shape[0] == 0: return None box = xywh2xyxy(pred[:, :4]) scores = pred[:, 4] * pred[:, 5:].max(1)[0] cls = pred[:, 5:].argmax(1) keep = torchvision.ops.nms(box, scores, iou_thres) return torch.cat([box[keep], scores[keep].unsqueeze(1), cls[keep].unsqueeze(1).float()], 1)xywh2xyxy是把中心点+宽高的格式转成左上角+右下角的格式,这是NMS需要的输入格式。这个转换看起来简单,但写错的人不少,尤其是x1 = x - w/2和x2 = x + w/2这两个公式,容易把w和h搞混。3.4 完整推理脚本的组装把预处理、推理、后处理串起来,就是一个完整的推理脚本。我习惯把它封装成一个类,方便复用:class Yolov5Detector: def __init__(self, weights, device='cuda', img_size=640, conf_thres=0.25, iou_thres=0.45): self.device = device self.img_size = img_size self.conf_thres = conf_thres self.iou_thres = iou_thres ckpt = torch.load(weights, map_location='cpu') self.model = ckpt['ema' if ckpt.get('ema') else 'model'].float() self.model.eval() self.model.to(device) self.names = ckpt['model'].names if hasattr(ckpt['model'], 'names') else None def detect(self, img_path): img, img0, ratio, pad = preprocess(img_path, self.img_size) img = img.to(self.device) with torch.no_grad(): pred = self.model(img)[0] det = non_max_suppression(pred, self.conf_thres, self.iou_thres) if det is None: return [] det[:, [0, 2]] -= pad[0] det[:, [1, 3]] -= pad[1] det[:, :4] /= ratio det[:, [0, 2]] = det[:, [0, 2]].clamp(0, img0.shape[1]) det[:, [1, 3]] = det[:, [1, 3]].clamp(0, img0.shape[0]) results = [] for *box, conf, cls in det: results.append({ 'box': [int(x) for x in box], 'conf': float(conf), 'cls': int(cls), 'name': self.names[int(cls)] if self.names else str(int(cls)) }) return results这个类可以直接拿去用,输入图片路径,输出检测结果列表。实际项目中,我会把img_path换成img数组,避免重复读盘。4. 性能优化与生产环境适配4.1 显存管理与批处理策略Pytorch部署Yolov5时,显存是稀缺资源。一张1080Ti有11GB显存,跑Yolov5s在640输入下单张推理大概占1.5GB,但如果你开batch=16,显存占用会飙升到8GB以上。所以batch size不能随便设,要根据显存和模型大小来算。一个粗略的估算公式:显存占用 ≈ 模型参数显存 + 激活值显存 × batch_size。Yolov5s的参数约7M,FP32下占28MB,FP16下占14MB,这部分很小。大头是激活值,也就是前向传播过程中每一层的输出。Yolov5s在640输入下,单张图的激活值约1.2GB(FP32),所以batch=8时约9.6GB,加上模型和其他开销,11GB显存基本就满了。实际部署时,我通常会用torch.cuda.memory_allocated()来监控显存,动态调整batch size。如果显存不够,优先考虑:1)切FP16,显存直接减半;2)降低输入尺寸,比如从640降到416,激活值按平方比例减少;3)用梯度累积模拟大batch(但推理时不需要梯度,所以这条不适用)。还有一个技巧:推理时用torch.inference_mode()替代torch.no_grad()。前者是Pytorch 1.9引入的,比后者更彻底,能进一步减少显存占用和提升速度。实测下来,差异在5%左右,但积少成多。4.2 多线程与异步推理生产环境里,推理服务通常要同时处理多个请求。如果每个请求都串行跑,吞吐量上不去。Pytorch本身不是线程安全的,但推理时如果模型是只读的,多线程调用model()是可以的。不过GIL的存在会让Python层面的多线程效果打折扣,所以更推荐用多进程或者异步IO + 单线程推理。我常用的方案是:用一个队列接收请求,一个后台线程专门跑推理,推理结果通过另一个队列返回。这样推理线程可以连续跑,不用等IO。如果请求量很大,就起多个进程,每个进程加载一份模型,用共享内存或消息队列分发请求。这个方案在4核CPU + 1张GPU的机器上,能跑到每秒50-80张图(Yolov5s,640输入,FP16)。注意:多进程加载模型时,每个进程都会占一份显存。如果显存不够,可以用torch.multiprocessing的spawn模式,配合torch.cuda.set_device()让每个进程用不同的GPU。但单GPU下,多进程反而会因为显存竞争变慢,这时候多线程 + 队列更合适。4.3 模型量化与加速的取舍Pytorch支持动态量化和静态量化,能把FP32模型转成INT8,显存占用减少75%,速度提升2-3倍。但Yolov5的量化不是开箱即用的,因为检测头里的某些操作(比如sigmoid、exp)对量化敏感,直接量化会导致精度大幅下降。我试过用Pytorch的torch.quantization.quantize_dynamic对Yolov5做动态量化,结果mAP掉了15个百分点,基本不可用。后来改用量化感知训练(QAT),在训练时插入伪量化节点,让模型适应量化误差,最终mAP只掉了2个百分点,速度提升了2.5倍。但QAT需要重新训练,成本较高,适合对速度要求极高的场景。如果不想重新训练,又想加速,可以考虑TorchScript。把模型转成TorchScript后,Pytorch会做图优化,去掉一些冗余操作,速度能提升10%-20%。转换方法很简单:scripted_model = torch.jit.script(model) scripted_model.save('yolov5s_scripted.pt')但TorchScript对动态控制流支持有限,Yolov5的某些版本可能转不过去。我实测Yolov5 v6.0可以转,v7.0有些版本会报错。转换前先用torch.jit.trace试一下,如果trace能过,script通常也能过。4.4 服务化部署的接口设计如果要把Yolov5的Pytorch推理做成一个服务,接口设计要考虑几点:输入格式、输出格式、错误处理、并发控制。输入格式我通常用JSON,包含image字段(base64编码的图片)和可选参数(conf_thres、iou_thres)。输出也是JSON,包含boxes、scores、classes三个数组。这样前端和后端都好处理。错误处理要覆盖:图片解码失败、模型推理异常、显存不足。显存不足时,可以捕获RuntimeError,然后清空torch.cuda.empty_cache(),再重试一次。如果还不行,就返回503,让客户端稍后重试。并发控制用信号量或者线程池。我一般用concurrent.futures.ThreadPoolExecutor,设置max_workers等于GPU能承受的并发数。每个worker里跑一次推理,结果通过Future返回。这样既能控制并发,又能复用线程。5. 常见问题排查与实战避坑指南5.1 加载权重报错:KeyError与Unexpected key这是最常见的问题,没有之一。报错信息通常是Missing key(s) in state_dict或者Unexpected key(s) in state_dict。原因一般是训练和部署的Yolov5代码版本不一致,导致模型定义里的层名变了。解决方法有两种:一是统一代码版本,把训练时的models目录复制到部署环境;二是手动映射键名,用collections.OrderedDict重新组织state_dict。第二种方法麻烦但灵活,适合无法统一版本的场景。具体操作:先打印训练时权重的键名和部署时模型的键名,对比差异,然后写一个映射函数。比如训练时叫model.0.conv.weight,部署时叫model.0.conv.0.weight,那就把前者映射到后者。这个工作很枯燥,但一次写好就能复用。5.2 推理结果为空或框位置错乱推理结果为空,通常是置信度阈值设太高,或者预处理时颜色通道搞错了。先检查conf_thres是不是0.25以上,再检查cv2.cvtColor有没有加。如果框位置错乱,比如框整体偏移或者大小不对,那多半是后处理的坐标映射出了问题。检查pad和ratio的计算是否正确,尤其是pad的左右和上下顺序。还有一个隐蔽的坑:letterbox的padding顺序。cv2.copyMakeBorder的参数顺序是top, bottom, left, right,而计算dw和dh时,dw是宽度的padding,对应left和right;dh是高度的padding,对应top和bottom。如果搞反了,框会整体偏移。我见过有人把dw和dh搞反,结果框全部偏到右下角。5.3 显存泄漏与OOM排查显存泄漏的表现是:跑一段时间后,显存占用越来越高,最终OOM。原因通常是:1)没有用torch.no_grad(),计算图没释放;2)中间变量(比如pred、det)一直持有引用,没被GC回收;3)多线程下,每个线程都创建了新的CUDA上下文。排查方法:用torch.cuda.memory_summary()看显存分配情况,找出哪个部分在增长。如果是计算图没释放,加torch.no_grad();如果是变量引用,手动del并调torch.cuda.empty_cache();如果是多线程上下文,改用单线程推理或者torch.multiprocessing。提示:torch.cuda.empty_cache()只能释放未被引用的显存,如果变量还被引用,调了也没用。所以关键是找到并删除引用。5.4 常见问题速查表问题现象可能原因排查方法解决方案加载权重报KeyError代码版本不一致对比键名统一版本或手动映射推理结果为空置信度阈值过高打印原始输出降低conf_thres框位置偏移padding顺序错误检查dw/dh修正copyMakeBorder参数显存持续增长计算图未释放memory_summary加no_grad,删引用推理速度慢未用FP16检查dtype转half,用inference_mode多线程结果乱序共享变量竞争加锁或改队列用队列串行化颜色通道错误BGR未转RGB检查cvtColor加BGR2RGB转换输入尺寸不对齐非32倍数检查img_size调整为32的倍数5.5 独家避坑经验第一个经验:永远保留一份“黄金输入”。就是一张固定的测试图,每次改完代码都跑一遍,看输出是否一致。这样能快速发现回归问题。我习惯把这张图的检测结果(框坐标、分数)存成JSON,改代码后对比差异,超过阈值就报警。第二个经验:预处理和后处理要成对测试。单独测预处理,看输出的tensor形状和数值范围;单独测后处理,用固定的tensor输入,看输出的框是否正确。这样能把问题定位到具体环节,而不是笼统地“推理不对”。第三个经验:GPU和CPU的结果要对比。有时候GPU上跑出来正常,CPU上跑出来不对,原因是某些操作在CPU和GPU上的实现有细微差异(比如插值算法)。如果部署环境可能没有GPU,一定要在CPU上验证一遍。第四个经验:模型切换要热更新。生产环境里,模型可能会更新。如果每次更新都重启服务,用户体验很差。我的做法是:把模型加载封装成一个函数,用文件监听或者API触发重新加载。加载新模型时,先加载到临时变量,验证通过后再替换旧模型,避免加载失败导致服务不可用。6. 从Pytorch部署到其他方案的过渡Pytorch部署跑通之后,你可能会遇到性能瓶颈,比如单张推理要50ms,但业务要求10ms。这时候就需要考虑其他部署方案了。Pytorch部署的最大价值,其实是作为基准:它告诉你模型在理想情况下的精度上限,后续任何加速方案(ONNX、TensorRT、OpenVINO)都要和它对比,确保精度损失在可接受范围内。从Pytorch到ONNX的转换,核心是torch.onnx.export。但Yolov5的导出有几个坑:1)输出需要是tuple而不是list;2)动态batch需要设置dynamic_axes;3)某些算子(比如torch.nn.functional.interpolate)在ONNX里的实现可能和Pytorch不同。我通常会在导出后,用onnxruntime跑一遍,和Pytorch的结果对比,确保差异小于1e-3。从ONNX到TensorRT,又是另一个坑。TensorRT对算子支持有限,Yolov5的某些版本需要插件(plugin)才能转。而且TensorRT的精度和Pytorch可能有差异,尤其是FP16模式下。我的建议是:先用Pytorch部署验证业务逻辑,再逐步替换推理后端。这样即使加速方案出问题,也有一个可用的兜底方案。最后分享一个我个人的习惯:把Pytorch推理脚本和模型权重一起打包成Docker镜像。这样部署时不用操心环境问题,直接docker run就能跑。镜像里只装必要的依赖(Pytorch、OpenCV、NumPy),体积能控制在2GB以内。如果显存紧张,还可以用nvidia-docker限制GPU内存,避免影响其他服务。