AI对抗训练折腾手记

去年做图像分类项目时踩过对抗样本的坑：用户随便调一下对比度，模型就把猫认成狗。下面记录后来用 FGSM、PGD、TRADES 做对抗训练的取舍和结果。

问题背景

去年我们在一个图像分类项目上遇到了这事。模型在验证集上表现很好，但上线后发现有个用户反馈：他上传的照片明明是"猫"，模型一直说"狗"。排查半天发现，用户用图片编辑器稍微调了下对比度——幅度肉眼都很难察觉，模型就崩了。

这就是对抗样本的问题。模型不是"学会了"，只是"背住了"训练数据。在训练分布内表现好，稍微偏离一点就出事。

这张图说明了一个很现实的问题：模型的"理解"和人的"理解"不在同一个维度。我们看的是图像内容，模型依赖的可能是某些高频特征的细微变化。

这次折腾的目标很简单：干净图和被篡改过的图都要尽量认对。

需求和限制

实际需求比理论上的"提升鲁棒性"要具体一些：

1. 模型本身是 ResNet-50，从头训练太贵，得用预训练模型微调
2. 训练数据是 ImageNet 子集，大约 10 万张图，有限算力下跑不完 90 轮
3. 对抗训练本身会降低干净样本的准确率，要找到一个平衡点
4. 上线时不能每次推理都做对抗攻击检测，推理开销不能太大

这些限制把目标压得很具体：有限算力下先让模型不那么脆，不追理论最优鲁棒性。

实现思路

对抗训练的思路也直白：训练阶段就把被人动过的样本混进来，别只喂干净数据。

生成对抗样本的算法不少，我这次试了 FGSM、PGD 和 TRADES，各有各的特点。

FGSM：最简单的入门选择

FGSM (Fast Gradient Sign Method) 是对抗训练的入门级算法，原理一句话能说清：沿着让损失增加最大的方向调整输入。

import torch
import torch.nn as nn

def fgsm_attack(model, x, y, epsilon=0.03):
    """
    x: 输入图像
    y: 真实标签
    epsilon: 扰动幅度
    """
    x_adv = x.clone().detach().requires_grad_(True)
    output = model(x_adv)
    loss = nn.CrossEntropyLoss()(output, y)
    loss.backward()

    # 梯度符号方向
    perturbation = epsilon * x_adv.grad.sign()

    # 限制在 epsilon 范围内
    x_adv = torch.clamp(x + perturbation, 0, 1)

    return x_adv

训练时把对抗样本和干净样本混在一起喂给模型：

def train_step(model, optimizer, x, y, epsilon=0.03):
    model.train()

    # 生成对抗样本
    x_adv = fgsm_attack(model, x, y, epsilon)

    # 混合训练
    mixed_x = torch.cat([x, x_adv], dim=0)
    mixed_y = torch.cat([y, y], dim=0)

    output = model(mixed_x)
    loss = nn.CrossEntropyLoss()(output, mixed_y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()

FGSM 的好处是快，生成一次对抗样本只需一次前向+反向。但它的问题也明显：单步扰动往往不是最强的攻击，模型可能学到"防御 FGSM"，但对更强的攻击没抵抗力。

实际跑下来，用 FGSM 做对抗训练后，模型对 FGSM 攻击的准确率从 15% 提升到 85%，但对 PGD 攻击的准确率还是只有 30% 左右。这就像练拳只练直拳，遇到对方组合拳还是不行。

PGD：更强的对抗性

PGD (Projected Gradient Descent) 可以看成是多步 FGSM。每次只走一小步，多次迭代，每步都把扰动限制在 epsilon 范围内。

def pgd_attack(model, x, y, epsilon=0.03, alpha=0.005, num_steps=10):
    """
    epsilon: 最大扰动幅度
    alpha: 每步步长
    num_steps: 迭代次数
    """
    x_adv = x.clone().detach() + torch.zeros_like(x).uniform_(-epsilon, epsilon)
    x_adv = torch.clamp(x_adv, 0, 1).requires_grad_(True)

    for _ in range(num_steps):
        output = model(x_adv)
        loss = nn.CrossEntropyLoss()(output, y)

        loss.backward()

        # 步长方向更新
        x_adv = x_adv + alpha * x_adv.grad.sign()

        # 投影到 epsilon 球内
        delta = torch.clamp(x_adv - x, -epsilon, epsilon)
        x_adv = torch.clamp(x + delta, 0, 1).detach().requires_grad_(True)

    return x_adv

PGD 生成的是更强的对抗样本，用 PGD 训练出来的模型鲁棒性也更好。但代价也明显：每张图要跑 10 次前向+反向，训练时间直接翻 10 倍。

我们的训练集有 10 万张图，用 PGD 做对抗训练，一轮要跑 100 万次前向+反向。单卡 3090 跑一轮要 6 小时，跑完 30 轮要一周。这对于一个小团队来说压力有点大。

TRADES：在鲁棒性和准确率之间找平衡

TRADES (TRadeoff between Accuracy and Robustness for Deep neural networks) 是 2019 年提出的算法：分类 loss 之外，还要拉齐对抗样本和干净样本的特征表示。

def trades_loss(model, x, y, x_adv, beta=6.0):
    """
    beta: 鲁棒性和准确率的权衡参数
    """
    # 自然样本的分类损失
    ce_loss = nn.CrossEntropyLoss()(model(x), y)

    # KL 散度，约束对抗样本和干净样本的预测分布
    kl_loss = nn.KLDivLoss(reduction='batchmean')(
        nn.LogSoftmax(dim=1)(model(x_adv)),
        nn.Softmax(dim=1)(model(x))
    )

    return ce_loss + beta * kl_loss

TRADES 的关键在于 beta 参数。beta 越大，模型越关注鲁棒性，但会牺牲干净样本的准确率。我们试了几个值：

最后选了 beta=6.0，虽然干净样本准确率掉了 7 个点，但鲁棒性提升了 16 个点，这个 trade-off 我们能接受。

踩坑记录

实践过程中遇到的几个问题，比预期要麻烦。

对抗样本生成导致显存爆炸

刚开始直接用 PGD 生成对抗样本，训练时 batch size 设为 128。结果显存直接溢出。一查才发现：PGD 每步都要保存梯度，10 步下来中间计算图都占着显存。

解决方法是每步都 .detach()，不用中间结果的梯度：

x_adv = x_adv.detach().requires_grad_(True)  # 每步都新建计算图

或者干脆先批量生成对抗样本，再训练。这样要存两张 batch 的图，但至少不会有多步梯度堆积。

训练时间太长

PGD 带来的训练时间开销让我们一度想放弃。后来用了两个办法缓解：

1. 减少 PGD 步数。从 10 步降到 7 步，鲁棒性掉得不多，但时间省了 30%。
2. 对抗训练只在最后 10 轮用。前 20 轮正常训练，最后 10 轮用对抗训练 fine-tune。这样总共训练时间只增加了 50%。

效果上，从头做对抗训练和最后 10 轮做对抗训练，最终鲁棒性差别不大。原因是模型在正常训练后期已经学会了主要特征，对抗训练更多是在"补边界情况"。

epsilon 太小鲁棒性不够，太大破坏语义

epsilon 控制扰动的最大幅度。我们试了不同值：

• epsilon=0.01：扰动很小，肉眼看不出来，但模型鲁棒性提升有限
• epsilon=0.03：扰动能看出一点痕迹，但不会影响人的判断，鲁棒性明显提升
• epsilon=0.07：扰动已经很明显，图片像被噪点覆盖，训练出来的模型对干净样本表现变差

最后选了 epsilon=0.03。这个值在实际场景中比较合理：真正的攻击者不会用太明显的扰动，否则用户一眼就能看出来。

最终结果

折腾了一圈，最终的效果：

我们最后选了 TRADES，因为它在干净样本准确率和鲁棒性之间找到了一个相对平衡点。PGD 鲁棒性最好，但干净样本准确率掉得有点多。FGSM 训练成本低，但防御能力不够。

上线后的效果也还行，用户反馈那些"稍微调了下对比度就认错"的情况少了很多。当然，对抗训练不是万能的，遇到更复杂的攻击还是可能失效。

结语

这次最后选了 TRADES（beta=6.0），干净准确率掉 7 个点，PGD-10 鲁棒性升 16 个点，团队能接受。PGD 鲁棒性最好但算力扛不住，FGSM 便宜但防不住强攻击。

对抗训练能缓解上线后"稍微调对比度就认错"这类问题，遇到更复杂的攻击还是会失效。算力、准确率、鲁棒性三者得一起估，没有免费午餐。

版权声明: 本文首发于 指尖魔法屋-AI对抗训练折腾手记（https://blog.thinkmoon.cn/post/329-ai-adversarial-training-fragile-robust-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！