AugMix数据增强实战：从原理到PyTorch实现

1次阅读

没有评论

共计 2296 个字符，预计需要花费 6 分钟才能阅读完成。

在计算机视觉任务中，数据增强是提升模型泛化能力的重要手段。传统的数据增强方法主要包括几何变换（如旋转、翻转、裁剪）和颜色变换（如亮度、对比度调整）。这些方法虽然简单易用，但也存在明显的局限性：

变换方式单一，缺乏组合性，难以模拟真实世界中的复杂变化
随机性有限，容易导致模型过拟合特定的增强模式
对对抗样本和自然干扰的鲁棒性提升有限

AugMix 通过分层随机混合基础增强操作，有效解决了这些问题。它不仅增加了数据多样性，还通过一致性损失保证了增强后的图像语义不变性。

AugMix 的核心思想可以分解为三个层次：

基础增强链 ：由 3 - 5 个基础增强操作（如旋转、平移、颜色变化等）随机组合而成
混合比例 ：使用 Dirichlet 分布（参数 α 控制）生成混合权重
随机权重 ：最终输出是原始图像与两个增强链的凸组合

数学表达为：

I_mix = w_orig * I + w_1 * chain_1(I) + w_2 * chain_2(I)

其中 w_orig + w_1 + w_2 = 1

AugMix 还引入了 JS 一致性损失（Jensen-Shannon Divergence），鼓励模型对原始图像和增强图像产生一致的预测：

L_consistency = JS(p(y|x), p(y|augmix(x)))

首先实现支持概率性执行的基础操作类：

import torch
import random
from PIL import ImageOps, ImageFilter

class AugmentOp:
    def __init__(self, op_name: str, prob: float = 0.5, magnitude: float = 1.0):
        self.prob = prob
        self.magnitude = magnitude
        self.op_name = op_name

    def __call__(self, img: torch.Tensor) -> torch.Tensor:
        if random.random() > self.prob:
            return img

        if self.op_name == 'rotate':
            angle = random.uniform(-30, 30) * self.magnitude
            return img.rotate(angle)
        elif self.op_name == 'translate':
            # 省略其他操作实现...
            pass

def augmix(image: torch.Tensor, ops_list: list, alpha: float = 1.0) -> torch.Tensor:
    """
    Args:
        image: Input tensor (C,H,W)
        ops_list: List of base augmentation operations
        alpha: Dirichlet distribution parameter
    """
    # 生成两个增强链
    chain1 = apply_augment_chain(image, ops_list)
    chain2 = apply_augment_chain(image, ops_list)

    # 采样混合权重
    m = torch.distributions.dirichlet.Dirichlet(torch.tensor([alpha, alpha, alpha]))
    weights = m.sample()

    # 凸组合
    mixed = weights[0] * image + weights[1] * chain1 + weights[2] * chain2
    return mixed

from torchvision import datasets, transforms

train_transform = transforms.Compose([transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    # AugMix 作为最后一个 transform
    lambda x: augmix(x, ops_list=[AugmentOp('rotate', 0.8),
        AugmentOp('color', 0.8),
        # 其他操作...
    ])
])

train_set = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)

在 CIFAR-10 上的测试结果（ResNet18 模型）：