【PyTorch】YOLOv3多目标检测实战：从数据准备到模型训练

1. 为什么选择YOLOv3进行多目标检测第一次接触目标检测任务时我被各种算法搞得眼花缭乱。Faster R-CNN、SSD、RetinaNet...直到遇到YOLOv3才发现这就是我想要的全能选手。YOLOv3最大的特点就是快真正做到了实时检测。我在一个停车场车辆检测项目里实测过在1080Ti显卡上能达到45FPS这比很多两阶段检测器快了近10倍。速度优势来源于它的单阶段检测架构。不同于那些先提候选框再分类的方法YOLOv3直接把图像划分成网格每个网格同时预测边界框和类别概率。这种端到端的设计特别适合需要实时处理的场景比如智能监控、自动驾驶这些领域。另一个让我惊喜的是它对小目标的检测能力。通过融合不同尺度的特征图就是著名的FPN结构YOLOv3可以同时捕捉大物体和小物体的特征。记得有次测试时它甚至检测出了200米外的一个行人而其他算法早就把这个人影忽略掉了。2. 搭建PyTorch开发环境工欲善其事必先利其器。配置环境时我推荐使用conda创建虚拟环境这样可以避免各种依赖冲突。我常用的配置是conda create -n yolov3 python3.8 conda activate yolov3 pip install torch1.10.0cu113 torchvision0.11.1cu113 -f https://download.pytorch.org/whl/torch_stable.html这里有个坑要注意PyTorch版本和CUDA版本的匹配。我有次用CUDA 10.2装了PyTorch 1.8结果训练时老是报莫名其妙的错误折腾半天才发现是版本不兼容。现在学乖了直接去官网查兼容矩阵。除了PyTorch还需要安装几个必备工具包pip install opencv-python matplotlib tqdm pillow验证环境是否正常import torch print(torch.__version__) # 应该输出1.10.0 print(torch.cuda.is_available()) # 必须返回True3. 准备COCO数据集实战COCO数据集是目标检测领域的标准答案包含80个常见物体类别。我第一次下载时被它的体积吓到了——超过18GB后来发现其实可以按需下载。最稳妥的下载方式是使用官方脚本mkdir -p data/coco cd data/coco wget https://pjreddie.com/media/files/val2014.zip unzip val2014.zip wget https://pjreddie.com/media/files/train2014.zip unzip train2014.zip数据集目录结构应该是这样的data/ └── coco/ ├── annotations/ # 标注文件 ├── train2014/ # 训练图片 └── val2014/ # 验证图片处理标注文件时我踩过一个坑COCO的原始标注是JSON格式而YOLOv3需要的是每张图片对应的txt文件。转换脚本可以这样写import json import os with open(annotations/instances_train2014.json) as f: data json.load(f) for img in data[images]: img_id img[id] img_w img[width] img_h img[height] anns [a for a in data[annotations] if a[image_id] img_id] with open(flabels/train2014/{img_id}.txt, w) as f: for ann in anns: x,y,w,h ann[bbox] x_center (x w/2) / img_w y_center (y h/2) / img_h w / img_w h / img_h f.write(f{ann[category_id]} {x_center} {y_center} {w} {h}\n)4. 构建自定义数据集加载器PyTorch的Dataset类让数据加载变得异常灵活。我设计CocoDataset时主要考虑三个需求支持数据增强处理不同尺寸的图片高效读取标注先看基础结构from torch.utils.data import Dataset import torchvision.transforms as T class CocoDataset(Dataset): def __init__(self, img_dir, label_dir, transformsNone): self.img_dir img_dir self.label_dir label_dir self.transforms transforms self.img_names os.listdir(img_dir) def __len__(self): return len(self.img_names) def __getitem__(self, idx): img_path os.path.join(self.img_dir, self.img_names[idx]) label_path os.path.join(self.label_dir, self.img_names[idx].replace(.jpg,.txt)) img Image.open(img_path).convert(RGB) # 读取标注 with open(label_path) as f: labels [list(map(float, line.split())) for line in f.read().splitlines()] if self.transforms: img, labels self.transforms(img, labels) return img, torch.tensor(labels)数据增强是提升模型泛化能力的关键。我常用的增强组合包括随机水平翻转p0.5色彩抖动随机裁剪Mosaic增强把4张图拼成1张def random_flip(img, labels): if random.random() 0.5: img TF.hflip(img) labels[:, 1] 1 - labels[:, 1] # 翻转x坐标 return img, labels5. 实现YOLOv3模型架构YOLOv3的骨干网络Darknet-53是个宝藏。我复现时发现它的残差连接设计比ResNet更高效。模型主要分三部分骨干网络特征提取class DarknetBlock(nn.Module): def __init__(self, in_channels): super().__init__() self.conv1 nn.Conv2d(in_channels, in_channels//2, 1) self.conv2 nn.Conv2d(in_channels//2, in_channels, 3, padding1) self.bn nn.BatchNorm2d(in_channels) def forward(self, x): residual x x F.leaky_relu(self.conv1(x), 0.1) x F.leaky_relu(self.bn(self.conv2(x)), 0.1) return x residual特征金字塔多尺度预测class FPN(nn.Module): def __init__(self): super().__init__() self.upsample nn.Upsample(scale_factor2, modenearest) def forward(self, x1, x2, x3): # x3是最大尺度的特征图 p3 self.conv3(x3) # 上采样并与x2融合 p2 self.conv2(x2 self.upsample(p3)) # 继续上采样 p1 self.conv1(x1 self.upsample(p2)) return p1, p2, p3检测头预测框和类别class DetectionHead(nn.Module): def __init__(self, anchors, num_classes): super().__init__() self.anchors anchors self.num_classes num_classes def forward(self, x): # x的形状: [batch, channels, height, width] batch, _, h, w x.shape # 调整维度用于预测 x x.view(batch, self.anchors, 5 self.num_classes, h, w) x x.permute(0,1,3,4,2) # 解码预测结果 boxes self._decode(x[..., :4]) conf torch.sigmoid(x[..., 4:5]) cls torch.sigmoid(x[..., 5:]) return torch.cat([boxes, conf, cls], dim-1)6. 训练技巧与调参经验训练YOLOv3就像驯服一匹野马需要耐心和技巧。我总结了几点关键经验学习率设置要用warmupdef warmup_lr(epoch): if epoch 3: return 0.001 * (epoch 1) / 3 elif epoch 30: return 0.01 elif epoch 60: return 0.001 else: return 0.0001损失函数是YOLOv3的灵魂包含三部分框坐标损失CIoU Loss置信度损失BCEWithLogitsLoss分类损失BCEWithLogitsLossclass YOLOLoss(nn.Module): def __init__(self): super().__init__() self.mse nn.MSELoss() self.bce nn.BCEWithLogitsLoss() def forward(self, pred, target): # 计算CIoU损失 iou bbox_iou(pred[..., :4], target[..., :4], CIoUTrue) box_loss (1 - iou).mean() # 置信度损失 obj_loss self.bce(pred[..., 4], target[..., 4]) # 分类损失 cls_loss self.bce(pred[..., 5:], target[..., 5:]) return box_loss obj_loss cls_loss数据增强方面Mosaic增强效果惊人。它能同时看到四张图的上下文def mosaic_augmentation(imgs, labels, size608): output_img np.zeros((size, size, 3)) output_labels [] # 随机选择四张图的拼接位置 xc, yc random.randint(size//2, size), random.randint(size//2, size) for i in range(4): img, anns imgs[i], labels[i] h, w img.shape[:2] if i 0: # 左上 x1a, y1a, x2a, y2a 0, 0, xc, yc x1b, y1b, x2b, y2b w - xc, h - yc, w, h elif i 1: # 右上 x1a, y1a, x2a, y2a xc, 0, size, yc x1b, y1b, x2b, y2b 0, h - yc, size - xc, h elif i 2: # 左下 x1a, y1a, x2a, y2a 0, yc, xc, size x1b, y1b, x2b, y2b w - xc, 0, w, size - yc else: # 右下 x1a, y1a, x2a, y2a xc, yc, size, size x1b, y1b, x2b, y2b 0, 0, size - xc, size - yc output_img[y1a:y2a, x1a:x2a] img[y1b:y2b, x1b:x2b] # 调整标注坐标 for ann in anns: ann[1] (ann[1] * w x1a - x1b) / size ann[2] (ann[2] * h y1a - y1b) / size output_labels.append(ann) return output_img, output_labels7. 模型评估与性能优化训练完成后评估指标不能只看mAP。我在实际项目中会关注不同尺度下的召回率每个类别的精确率推理速度FPS评估代码示例def evaluate(model, dataloader, iou_thresh0.5): model.eval() stats [] for imgs, targets in tqdm(dataloader): imgs imgs.to(device) with torch.no_grad(): outputs model(imgs) # 非极大值抑制 outputs non_max_suppression(outputs) for i, (pred, target) in enumerate(zip(outputs, targets)): # 计算TP, FP, FN ... # 计算mAP ap compute_ap(stats) return ap性能优化方面我常用的技巧包括使用混合精度训练Amp梯度累积适合小batch_size情况模型剪枝去除冗余通道from torch.cuda.amp import autocast autocast() def train_step(imgs, targets): imgs imgs.to(device) targets targets.to(device) with autocast(): outputs model(imgs) loss criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()8. 实际部署中的坑与解决方案把模型部署到生产环境时遇到的第一个问题就是模型体积。原始YOLOv3有200多MB经过以下优化可以缩小到30MB模型量化model model.cpu() quantized_model torch.quantization.quantize_dynamic( model, {nn.Conv2d, nn.Linear}, dtypetorch.qint8 ) torch.jit.save(torch.jit.script(quantized_model), yolov3_quantized.pt)ONNX导出dummy_input torch.randn(1, 3, 416, 416) torch.onnx.export(model, dummy_input, yolov3.onnx, opset_version11, do_constant_foldingTrue)另一个常见问题是推理速度不稳定。通过以下方法可以优化固定输入尺寸避免动态resize使用TensorRT加速启用CUDA Graph# TensorRT优化示例 import tensorrt as trt logger trt.Logger(trt.Logger.INFO) builder trt.Builder(logger) network builder.create_network() parser trt.OnnxParser(network, logger) with open(yolov3.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.max_workspace_size 1 30 engine builder.build_engine(network, config)最后提醒一个容易忽视的问题内存泄漏。长期运行的推理服务需要定期检查import gc def inference(img): with torch.no_grad(): output model(img) # 手动释放内存 del output gc.collect() torch.cuda.empty_cache()