保姆级教程：用PyTorch从零复现YOLOv4（附完整代码与Mosaic数据增强实现）

从零构建YOLOv4代码级实现与核心模块解析1. 环境配置与工具准备在开始复现YOLOv4之前我们需要搭建一个高效的开发环境。推荐使用Python 3.8和PyTorch 1.7的组合这是目前最稳定的深度学习开发环境之一。首先安装必要的依赖库pip install torch1.7.1 torchvision0.8.2 pip install opencv-python numpy tqdm matplotlib对于GPU加速建议使用CUDA 11.0配合cuDNN 8.0.5。可以通过以下命令验证PyTorch是否正确识别了GPUimport torch print(torch.cuda.is_available()) # 应输出True print(torch.cuda.device_count()) # 显示可用GPU数量提示如果使用Docker环境推荐使用官方PyTorch镜像作为基础镜像可以避免很多环境配置问题。2. 数据准备与Mosaic增强实现YOLOv4最显著的特点之一是其创新的Mosaic数据增强技术。这种技术将四张训练图像拼接成一张极大地丰富了训练数据的多样性。2.1 基础数据加载器首先实现一个基础的COCO格式数据加载器import cv2 import numpy as np from torch.utils.data import Dataset class YOLODataset(Dataset): def __init__(self, img_dir, label_dir, img_size416): self.img_dir img_dir self.label_dir label_dir self.img_size img_size self.img_files [f for f in os.listdir(img_dir) if f.endswith(.jpg)] def __getitem__(self, index): img_path os.path.join(self.img_dir, self.img_files[index]) label_path os.path.join(self.label_dir, self.img_files[index].replace(.jpg, .txt)) img cv2.imread(img_path) img cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 更多预处理代码... return img, label2.2 Mosaic增强实现下面是Mosaic增强的核心代码实现def mosaic_augmentation(self, index): # 随机选择4张图片的索引 indices [index] [random.randint(0, len(self) - 1) for _ in range(3)] # 创建拼接后的画布 mosaic_img np.zeros((self.img_size * 2, self.img_size * 2, 3), dtypenp.uint8) mosaic_labels [] # 四等分位置 positions [ (0, 0), # 左上 (self.img_size, 0), # 右上 (0, self.img_size), # 左下 (self.img_size, self.img_size) # 右下 ] for i, (x, y) in zip(indices, positions): img, label self.load_single_item(i) h, w img.shape[:2] # 随机缩放 scale random.uniform(0.5, 1.5) img cv2.resize(img, (int(w * scale), int(h * scale))) # 随机裁剪并放置到对应位置 crop_x random.randint(0, max(0, img.shape[1] - self.img_size)) crop_y random.randint(0, max(0, img.shape[0] - self.img_size)) img_crop img[crop_y:crop_yself.img_size, crop_x:crop_xself.img_size] mosaic_img[y:yself.img_size, x:xself.img_size] img_crop # 调整标签坐标 for obj in label: x1, y1, x2, y2, cls obj # 坐标转换代码... mosaic_labels.append(adjusted_obj) return mosaic_img, mosaic_labels注意Mosaic增强会显著增加GPU显存使用量建议适当减小batch size。3. CSPDarknet53骨干网络实现YOLOv4采用了CSPDarknet53作为骨干网络这是其高效性的关键所在。下面我们分模块实现这个网络。3.1 基础构建块首先实现CSPDarknet的基本构建块import torch.nn as nn class ConvBNMish(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride1): super().__init__() padding (kernel_size - 1) // 2 self.conv nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, biasFalse) self.bn nn.BatchNorm2d(out_channels) self.mish nn.Mish() def forward(self, x): return self.mish(self.bn(self.conv(x))) class ResidualBlock(nn.Module): def __init__(self, channels): super().__init__() self.conv1 ConvBNMish(channels, channels//2, 1) self.conv2 ConvBNMish(channels//2, channels, 3) def forward(self, x): residual x out self.conv1(x) out self.conv2(out) out residual return out3.2 CSP模块实现CSP(Cross Stage Partial)结构是CSPDarknet的核心创新class CSPBlock(nn.Module): def __init__(self, in_channels, out_channels, num_blocks): super().__init__() self.downsample ConvBNMish(in_channels, out_channels, 3, stride2) # 将通道分成两部分 mid_channels out_channels // 2 self.conv1 ConvBNMish(out_channels, mid_channels, 1) self.conv2 ConvBNMish(out_channels, mid_channels, 1) # 残差块部分 self.blocks nn.Sequential( *[ResidualBlock(mid_channels) for _ in range(num_blocks)] ) self.conv3 ConvBNMish(mid_channels, mid_channels, 1) self.conv4 ConvBNMish(2 * mid_channels, out_channels, 1) def forward(self, x): x self.downsample(x) # 分割路径 x1 self.conv1(x) x2 self.conv2(x) # 残差路径 x2 self.blocks(x2) x2 self.conv3(x2) # 合并路径 out torch.cat([x1, x2], dim1) out self.conv4(out) return out3.3 完整CSPDarknet53架构基于上述模块我们可以构建完整的CSPDarknet53class CSPDarknet53(nn.Module): def __init__(self): super().__init__() self.conv1 ConvBNMish(3, 32, 3) self.stage1 CSPBlock(32, 64, 1) self.stage2 CSPBlock(64, 128, 2) self.stage3 CSPBlock(128, 256, 8) self.stage4 CSPBlock(256, 512, 8) self.stage5 CSPBlock(512, 1024, 4) def forward(self, x): x self.conv1(x) x self.stage1(x) x self.stage2(x) out3 self.stage3(x) # 输出1/8尺度特征 out4 self.stage4(out3) # 输出1/16尺度特征 out5 self.stage5(out4) # 输出1/32尺度特征 return out3, out4, out54. 特征融合网络实现YOLOv4采用了SPP和PAN结构进行多尺度特征融合这是其高性能的关键之一。4.1 SPP(空间金字塔池化)模块class SPP(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() mid_channels in_channels // 2 self.conv1 ConvBNMish(in_channels, mid_channels, 1) self.conv2 ConvBNMish(mid_channels * 4, out_channels, 1) # 不同尺度的最大池化 self.pool5 nn.MaxPool2d(5, stride1, padding2) self.pool9 nn.MaxPool2d(9, stride1, padding4) self.pool13 nn.MaxPool2d(13, stride1, padding6) def forward(self, x): x self.conv1(x) pool1 x pool5 self.pool5(x) pool9 self.pool9(x) pool13 self.pool13(x) out torch.cat([pool1, pool5, pool9, pool13], dim1) out self.conv2(out) return out4.2 PAN(路径聚合网络)实现class PAN(nn.Module): def __init__(self, channels_list): super().__init__() # 上采样路径(自底向上) self.upsample nn.Upsample(scale_factor2, modenearest) # 定义多个卷积层 self.conv1 ConvBNMish(channels_list[0], channels_list[1], 1) self.conv2 ConvBNMish(channels_list[1], channels_list[1], 3) self.conv3 ConvBNMish(channels_list[1], channels_list[2], 1) self.conv4 ConvBNMish(channels_list[2], channels_list[2], 3) # 下采样路径(自顶向下) self.downsample ConvBNMish(channels_list[2], channels_list[1], 3, stride2) def forward(self, features): # features包含三个尺度的特征图 x_low, x_mid, x_high features # 上采样路径 up1 self.conv1(x_high) up1 self.upsample(up1) up1 torch.cat([up1, x_mid], dim1) up1 self.conv2(up1) up2 self.conv3(up1) up2 self.upsample(up2) up2 torch.cat([up2, x_low], dim1) up2 self.conv4(up2) # 下采样路径 down1 self.downsample(up2) down1 torch.cat([down1, up1], dim1) down1 self.conv2(down1) down2 self.downsample(down1) down2 torch.cat([down2, x_high], dim1) down2 self.conv2(down2) return up2, down1, down25. YOLOv4头部与损失函数5.1 YOLO头部实现YOLOv4沿用了YOLOv3的头部设计但做了一些改进class YOLOHead(nn.Module): def __init__(self, in_channels, num_classes, anchors): super().__init__() self.num_classes num_classes self.num_anchors len(anchors) # 预测层 self.conv1 ConvBNMish(in_channels, in_channels*2, 3) self.conv2 nn.Conv2d(in_channels*2, self.num_anchors*(5 num_classes), 1) # 初始化参数 nn.init.normal_(self.conv2.weight, mean0, std0.01) nn.init.constant_(self.conv2.bias, 0) # 锚框 self.register_buffer(anchors, torch.tensor(anchors).float()) def forward(self, x): batch_size x.size(0) grid_size x.size(2) # 特征提取 out self.conv1(x) out self.conv2(out) # 调整形状 out out.view(batch_size, self.num_anchors, 5 self.num_classes, grid_size, grid_size) out out.permute(0, 1, 3, 4, 2).contiguous() return out5.2 CIoU损失函数实现YOLOv4采用了CIoU(Complete Intersection over Union)损失函数class CIOULoss(nn.Module): def __init__(self, reductionmean): super().__init__() self.reduction reduction def forward(self, pred, target): # 预测框和目标框坐标 pred_left pred[..., 0] - pred[..., 2] / 2 pred_top pred[..., 1] - pred[..., 3] / 2 pred_right pred[..., 0] pred[..., 2] / 2 pred_bottom pred[..., 1] pred[..., 3] / 2 target_left target[..., 0] - target[..., 2] / 2 target_top target[..., 1] - target[..., 3] / 2 target_right target[..., 0] target[..., 2] / 2 target_bottom target[..., 1] target[..., 3] / 2 # 计算交集面积 intersect_left torch.max(pred_left, target_left) intersect_top torch.max(pred_top, target_top) intersect_right torch.min(pred_right, target_right) intersect_bottom torch.min(pred_bottom, target_bottom) intersect_area (intersect_right - intersect_left).clamp(min0) * \ (intersect_bottom - intersect_top).clamp(min0) # 计算并集面积 pred_area (pred_right - pred_left) * (pred_bottom - pred_top) target_area (target_right - target_left) * (target_bottom - target_top) union_area pred_area target_area - intersect_area # IoU计算 iou intersect_area / (union_area 1e-7) # 中心点距离 center_dist torch.pow(pred[..., :2] - target[..., :2], 2).sum(dim-1) # 最小包围框对角线距离 enclose_left torch.min(pred_left, target_left) enclose_top torch.min(pred_top, target_top) enclose_right torch.max(pred_right, target_right) enclose_bottom torch.max(pred_bottom, target_bottom) enclose_diag torch.pow(enclose_right - enclose_left, 2) \ torch.pow(enclose_bottom - enclose_top, 2) # CIoU计算 v (4 / (math.pi ** 2)) * torch.pow( torch.atan(target[..., 2] / (target[..., 3] 1e-7)) - torch.atan(pred[..., 2] / (pred[..., 3] 1e-7)), 2) alpha v / (1 - iou v 1e-7) ciou iou - (center_dist / (enclose_diag 1e-7) alpha * v) loss 1 - ciou if self.reduction mean: return loss.mean() elif self.reduction sum: return loss.sum() else: return loss6. 训练策略与技巧6.1 自对抗训练(SAT)YOLOv4引入了自对抗训练(Self-Adversarial Training)来提升模型鲁棒性def self_adversarial_training(model, images, targets, optimizer): # 第一阶段冻结权重生成对抗样本 for param in model.parameters(): param.requires_grad False # 添加随机噪声 noise torch.randn_like(images) * 0.1 noisy_images images noise noisy_images torch.clamp(noisy_images, 0, 1) # 计算损失并反向传播 predictions model(noisy_images) loss compute_loss(predictions, targets) loss.backward() # 获取梯度并生成对抗样本 grad images.grad.data adv_images images 0.3 * torch.sign(grad) adv_images torch.clamp(adv_images, 0, 1) # 第二阶段解冻权重正常训练 for param in model.parameters(): param.requires_grad True # 使用对抗样本训练 optimizer.zero_grad() predictions model(adv_images) loss compute_loss(predictions, targets) loss.backward() optimizer.step() return loss.item()6.2 DropBlock正则化YOLOv4使用DropBlock替代传统的Dropoutclass DropBlock(nn.Module): def __init__(self, block_size7, keep_prob0.9): super().__init__() self.block_size block_size self.keep_prob keep_prob def forward(self, x): if not self.training or self.keep_prob 1: return x # 计算gamma值 gamma (1 - self.keep_prob) / (self.block_size ** 2) # 创建掩码 mask_shape (x.shape[0], x.shape[1], x.shape[2] - self.block_size 1, x.shape[3] - self.block_size 1) mask torch.bernoulli(torch.full(mask_shape, gamma, devicex.device)) # 扩展掩码 mask F.pad(mask, [self.block_size // 2] * 4, value0) mask F.max_pool2d(mask, stride(1, 1), kernel_size(self.block_size, self.block_size), paddingself.block_size // 2) mask 1 - mask # 反转掩码 # 归一化以保持期望值不变 normalize_scale mask.numel() / (1e-7 mask.sum()) out x * mask * normalize_scale return out6.3 多尺度训练实现YOLOv4在训练过程中动态调整输入图像尺寸def random_resize(img, target_size): # 随机选择缩放比例 scale random.uniform(0.5, 1.5) new_size int(target_size * scale) # 调整图像大小 h, w img.shape[:2] ratio min(new_size / h, new_size / w) new_h, new_w int(h * ratio), int(w * ratio) img cv2.resize(img, (new_w, new_h)) # 填充到目标尺寸 pad_h target_size - new_h pad_w target_size - new_w top pad_h // 2 bottom pad_h - top left pad_w // 2 right pad_w - left img cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value(114, 114, 114)) return img7. 模型评估与推理优化7.1 DIoU-NMS实现YOLOv4采用了DIoU-NMS来改进后处理def diou_nms(boxes, scores, iou_threshold0.5): # 按置信度排序 order scores.argsort(descendingTrue) keep [] while order.size(0) 0: # 保留当前最高分框 i order[0] keep.append(i) if order.size(0) 1: break # 计算DIoU inter_left torch.max(boxes[i, 0], boxes[order[1:], 0]) inter_top torch.max(boxes[i, 1], boxes[order[1:], 1]) inter_right torch.min(boxes[i, 2], boxes[order[1:], 2]) inter_bottom torch.min(boxes[i, 3], boxes[order[1:], 3]) inter_area (inter_right - inter_left).clamp(min0) * \ (inter_bottom - inter_top).clamp(min0) area_i (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) area_j (boxes[order[1:], 2] - boxes[order[1:], 0]) * \ (boxes[order[1:], 3] - boxes[order[1:], 1]) union_area area_i area_j - inter_area iou inter_area / union_area # 计算中心点距离 center_i torch.stack([(boxes[i, 0] boxes[i, 2]) / 2, (boxes[i, 1] boxes[i, 3]) / 2]) center_j torch.stack([(boxes[order[1:], 0] boxes[order[1:], 2]) / 2, (boxes[order[1:], 1] boxes[order[1:], 3]) / 2], dim1) center_dist torch.sum(torch.pow(center_i - center_j, 2), dim0) # 计算最小包围框对角线距离 enclose_left torch.min(boxes[i, 0], boxes[order[1:], 0]) enclose_top torch.min(boxes[i, 1], boxes[order[1:], 1]) enclose_right torch.max(boxes[i, 2], boxes[order[1:], 2]) enclose_bottom torch.max(boxes[i, 3], boxes[order[1:], 3]) enclose_diag torch.pow(enclose_right - enclose_left, 2) \ torch.pow(enclose_bottom - enclose_top, 2) # DIoU计算 diou iou - center_dist / (enclose_diag 1e-7) # 保留DIoU低于阈值的框 mask diou iou_threshold order order[1:][mask] return torch.tensor(keep, dtypetorch.long)7.2 推理优化技巧在实际部署中我们可以采用多种优化技术def optimize_model_for_inference(model, img_size416): # 转换为评估模式 model.eval() # 模型量化 quantized_model torch.quantization.quantize_dynamic( model, {torch.nn.Linear, torch.nn.Conv2d}, dtypetorch.qint8 ) # 脚本化 scripted_model torch.jit.script(quantized_model) # 预热 dummy_input torch.randn(1, 3, img_size, img_size) for _ in range(10): _ scripted_model(dummy_input) return scripted_model8. 完整训练流程最后我们将所有组件整合成一个完整的训练流程def train_yolov4(): # 初始化模型 backbone CSPDarknet53() model YOLOv4(backbone, num_classes80).cuda() # 数据加载 train_dataset YOLODataset(train_img_dir, train_label_dir) train_loader DataLoader(train_dataset, batch_size16, shuffleTrue) # 优化器 optimizer torch.optim.AdamW(model.parameters(), lr1e-4, weight_decay5e-4) # 学习率调度 scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max100) # 训练循环 for epoch in range(300): model.train() for batch_idx, (images, targets) in enumerate(train_loader): images images.cuda() targets [t.cuda() for t in targets] # Mosaic增强 if random.random() 0.75: # 75%概率使用Mosaic images, targets mosaic_augmentation(images, targets) # 自对抗训练 if epoch 10 and random.random() 0.5: # 10个epoch后开始SAT loss self_adversarial_training(model, images, targets, optimizer) else: # 正常训练 optimizer.zero_grad() outputs model(images) loss compute_loss(outputs, targets) loss.backward() optimizer.step() # 打印训练信息 if batch_idx % 50 0: print(fEpoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item()}) # 更新学习率 scheduler.step() # 验证评估 if epoch % 5 0: evaluate(model, val_loader) # 保存模型 if epoch % 10 0: torch.save(model.state_dict(), fyolov4_epoch_{epoch}.pth)在实际项目中我发现Mosaic增强和自对抗训练的结合使用可以显著提升模型在小目标检测上的性能特别是在复杂场景中。通过调整Mosaic的使用概率(通常在50%-75%之间)可以在训练效率和模型性能之间取得良好平衡。