4. 数据加载与预处理：PyTorch Dataset与DataLoader、数据增强(Mosaic, MixUp, HSV)、归一化与Resize策略

数据加载和预处理，说白了就是给模型喂饭。饭喂得好不好，直接决定模型能不能长大个儿。我刚开始做YOLO项目时，总觉得数据加载就是读读图片、改改尺寸，结果训练出来的模型泛化能力差得离谱。后来才明白，这一块才是真正的技术活。

核心观点：数据预处理不是简单的"读图-改尺寸-送网络"，而是整个训练流程中最能体现工程经验的部分。好的预处理策略，能让模型收敛速度提升30%以上。

4.1 PyTorch Dataset与DataLoader：数据管道的基石

PyTorch的Dataset和DataLoader，我习惯把它们比作"数据食堂"和"打饭阿姨"。Dataset负责准备食材，DataLoader负责按需打饭。

4.1.1 自定义Dataset类

写一个YOLO专用的Dataset类，其实没那么复杂。核心就三个方法：__init__、__len__、__getitem__。

class YOLODataset(Dataset):
    def __init__(self, img_dir, label_dir, transforms=None):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.transforms = transforms
        self.img_files = [f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.png'))]
    
    def __len__(self):
        return len(self.img_files)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_files[idx])
        label_path = os.path.join(self.label_dir, 
                                  self.img_files[idx].replace('.jpg', '.txt'))
        
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        boxes = []
        if os.path.exists(label_path):
            with open(label_path, 'r') as f:
                for line in f.readlines():
                    class_id, x_center, y_center, w, h = map(float, line.strip().split())
                    boxes.append([class_id, x_center, y_center, w, h])
        
        if self.transforms:
            image, boxes = self.transforms(image, boxes)
        
        return image, torch.tensor(boxes, dtype=torch.float32)

嗯，这里要注意一点：YOLO的标签格式是归一化后的坐标，范围在0到1之间。我曾经踩过这个坑，直接把像素坐标送进去，结果模型训练了三天，loss死活不降。

4.1.2 DataLoader的配置技巧

DataLoader的配置，说白了就是调几个参数。但参数调得好，训练效率能翻倍。

参数 推荐值 说明
batch_size 16-64 取决于GPU显存，YOLO一般用16或32
num_workers 4-8 CPU核心数的一半，太多反而会慢
pin_memory True 加速GPU数据传输，前提是用了CUDA
shuffle True 训练时打乱，验证时不用
drop_last True 防止最后一个batch太小影响BN层

我的经验：num_workers设成4通常是最稳的。设太大反而会因为CPU切换开销导致IO变慢。我曾经试过设成16，结果数据加载比训练还慢。

4.2 数据增强：让模型见过世面

数据增强，说白了就是"无中生有"。一张图片通过变换，能变出十张、百张。模型见过的花样多了，自然就不容易过拟合。

4.2.1 Mosaic增强：YOLOv4的杀手锏

Mosaic增强，就是把四张图拼成一张。这样做的好处是：

  • 一张图里能看到多个场景，模型学得更快
  • 小目标被放大的概率增加，对小目标检测友好
  • 相当于变相增大了batch size
def mosaic_augmentation(image1, boxes1, image2, boxes2, 
                        image3, boxes3, image4, boxes4):
    # 随机选择拼接中心点
    cx, cy = random.randint(320, 960), random.randint(320, 960)
    
    # 创建画布
    mosaic_img = np.zeros((1280, 1280, 3), dtype=np.uint8)
    mosaic_boxes = []
    
    # 将四张图拼接到对应象限
    # 左上
    mosaic_img[0:cy, 0:cx] = cv2.resize(image1, (cx, cy))
    # 右上
    mosaic_img[0:cy, cx:1280] = cv2.resize(image2, (1280-cx, cy))
    # 左下
    mosaic_img[cy:1280, 0:cx] = cv2.resize(image3, (cx, 1280-cy))
    # 右下
    mosaic_img[cy:1280, cx:1280] = cv2.resize(image4, (1280-cx, 1280-cy))
    
    # 调整边界框坐标
    # ... (坐标变换逻辑)
    
    return mosaic_img, mosaic_boxes

注意：Mosaic增强会改变图片的统计分布。如果训练集和验证集分布差异太大，模型在验证集上可能会崩。我建议训练到后期时，逐渐降低Mosaic的概率。

4.2.2 MixUp增强：让模型学会"模糊"

MixUp，说白了就是把两张图按比例混合。比如0.7的猫加上0.3的狗，标签也按同样的比例混合。

def mixup_augmentation(image1, boxes1, image2, boxes2, alpha=0.2):
    lam = np.random.beta(alpha, alpha)
    
    mixed_image = lam * image1 + (1 - lam) * image2
    
    # 标签也要混合
    mixed_boxes = torch.cat([
        boxes1 * lam, 
        boxes2 * (1 - lam)
    ], dim=0)
    
    return mixed_image, mixed_boxes

我个人习惯在训练的前半段用MixUp，后半段关掉。因为模型前期需要"见多识广"，后期需要"精准打击"。

4.2.3 HSV增强：颜色不变，内容变

HSV增强，就是调整色调(Hue)、饱和度(Saturation)、明度(Value)。这样做的好处是让模型对光照变化不敏感。

def hsv_augmentation(image, h_gain=0.015, s_gain=0.7, v_gain=0.4):
    # 转换为HSV空间
    hsv = cv2.cvtColor(image, cv2.COLOR_RGB2HSV).astype(np.float32)
    
    # 随机调整
    hsv[:, :, 0] += random.uniform(-h_gain, h_gain) * 180
    hsv[:, :, 1] *= random.uniform(1 - s_gain, 1 + s_gain)
    hsv[:, :, 2] *= random.uniform(1 - v_gain, 1 + v_gain)
    
    # 裁剪到有效范围
    hsv = np.clip(hsv, 0, 255).astype(np.uint8)
    
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)

避坑指南：我曾经把H通道的增益设得太大，结果红色和蓝色混在一起，模型把红灯认成了绿灯。建议H增益不要超过0.02。

4.3 归一化与Resize策略

归一化和Resize，说白了就是让数据"标准化"。模型不喜欢"惊喜"，它希望看到的数据都在一个尺度上。

4.3.1 归一化：把数据拉到同一个起跑线

YOLO常用的归一化方式有两种：

方式 公式 适用场景
Min-Max归一化 (x - min) / (max - min) 像素值固定范围(0-255)
Z-score归一化 (x - mean) / std 数据分布已知的情况

YOLO官方用的是Min-Max归一化，把像素值从[0, 255]缩放到[0, 1]。然后减均值除方差，得到[-1, 1]的范围。

def normalize_image(image):
    # 先缩放到[0, 1]
    image = image / 255.0
    
    # 再标准化
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    
    image = (image - mean) / std
    
    return image

4.3.2 Resize策略：保持长宽比还是直接拉伸？

这个问题，我当年纠结了很久。直接拉伸会改变目标形状，保持长宽比又会引入黑边。

YOLO的做法是：保持长宽比，用黑边填充。这样做的好处是：

  • 目标形状不变，模型学到的特征更准确
  • 黑边部分不参与计算，不影响梯度
  • 可以统一输入尺寸，方便batch处理
def letterbox_resize(image, target_size=640):
    h, w = image.shape[:2]
    scale = min(target_size / h, target_size / w)
    
    new_h, new_w = int(h * scale), int(w * scale)
    resized = cv2.resize(image, (new_w, new_h))
    
    # 创建画布并填充黑边
    canvas = np.full((target_size, target_size, 3), 114, dtype=np.uint8)
    
    # 计算偏移量
    dw = (target_size - new_w) // 2
    dh = (target_size - new_h) // 2
    
    canvas[dh:dh+new_h, dw:dw+new_w] = resized
    
    return canvas, (scale, dw, dh)

注意：Resize后，边界框坐标也要跟着变换。我见过有人只resize图片不调整坐标，结果模型预测的框全偏了。

4.4 数据预处理流程图

下面这张图，是我自己总结的数据预处理流程。你想想看，从原始图片到模型输入，中间要经过多少道工序。

YOLO数据预处理流程图 原始图片 + 标签 YOLODataset Mosaic增强 MixUp增强 HSV增强 Letterbox Resize 归一化 [0,1] → [-1,1] 模型输入 Tensor 关键说明 • 数据增强随机组合使用 • 训练后期降低增强强度 • 验证集不做数据增强 • 边界框坐标同步变换

4.5 完整的数据加载流程

把上面这些串起来，就是一个完整的数据加载流程。我习惯这样组织代码：

def get_dataloader(img_dir, label_dir, batch_size=16, 
                   augment=True, mosaic_prob=0.5):
    # 定义数据增强
    transforms = []
    
    if augment:
        transforms.append(RandomHorizontalFlip(p=0.5))
        transforms.append(HSVAugmentation())
        
        if random.random() < mosaic_prob:
            transforms.append(MosaicAugmentation())
        
        transforms.append(MixUpAugmentation(alpha=0.2))
    
    # Resize和归一化
    transforms.append(LetterboxResize(target_size=640))
    transforms.append(Normalize())
    
    # 创建Dataset
    dataset = YOLODataset(
        img_dir=img_dir,
        label_dir=label_dir,
        transforms=Compose(transforms)
    )
    
    # 创建DataLoader
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        drop_last=True
    )
    
    return dataloader

我的习惯：训练时用Mosaic+MixUp+HSV的组合拳，验证时只用Resize和归一化。这样模型在训练时"见多识广"，验证时"实事求是"。

数据加载和预处理这块，说白了就是"磨刀不误砍柴工"。花时间把数据管道搭好，后面训练起来会顺畅很多。我见过太多人急着训练，结果数据加载成了瓶颈，GPU利用率不到50%。

← 上一章 📖 返回目录 下一章 →