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.1cu117这个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_locationcpu) 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_shape640, 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, interpolationcv2.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, valuecolor) return img, r, (dw, dh) def preprocess(img_path, img_size640): 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然后按batch8或16来跑吞吐量比逐张推理高好几倍。第四model.eval()不能忘。虽然Yolov5的推理模式下BN层和Dropout的行为和训练模式不同但如果不调eval()BN层会用batch统计量而不是滑动平均结果会不稳定。这个坑在单张推理时可能不明显但batch推理时差异会放大。3.3 后处理从25200个预测框到最终结果后处理是Yolov5部署里最复杂的部分。模型输出[1, 25200, 85]这25200个框里绝大多数是背景或低置信度的需要经过置信度过滤、NMS、坐标映射三步才能得到最终结果。置信度过滤很简单conf obj_conf * cls_confYolov5的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_thres0.25, iou_thres0.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, devicecuda, img_size640, conf_thres0.25, iou_thres0.45): self.device device self.img_size img_size self.conf_thres conf_thres self.iou_thres iou_thres ckpt torch.load(weights, map_locationcpu) 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但如果你开batch16显存占用会飙升到8GB以上。所以batch size不能随便设要根据显存和模型大小来算。一个粗略的估算公式显存占用 ≈ 模型参数显存 激活值显存 × batch_size。Yolov5s的参数约7MFP32下占28MBFP16下占14MB这部分很小。大头是激活值也就是前向传播过程中每一层的输出。Yolov5s在640输入下单张图的激活值约1.2GBFP32所以batch8时约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张图Yolov5s640输入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和rightdh是高度的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而不是list2动态batch需要设置dynamic_axes3某些算子比如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内存避免影响其他服务。