竞赛trick-AWP对抗训练的即插即用实现

1、对抗训练的简单概括

对抗训练就是使用对抗样本去训练模型，从而通过对原始训练数据添加噪声便得到了对抗样本。在竞赛中常在BERT的embedding阶段进行扰动，常使用的对抗训练有pgd,fgm,freelb等。本文主要记录awp对抗训练的即插即用实现。（缘故：在刚结束的腾讯微信大数据挑战赛-多模态短视频分类竞赛中有大幅度提升）。

2、实现

2.1 awp实现

class AWP:
    def __init__(self, model, optimizer, adv_param="weight", adv_lr=1, adv_eps=0.0001):
        self.model = model
        self.optimizer = optimizer
        self.adv_param = adv_param
        self.adv_lr = adv_lr
        self.adv_eps = adv_eps
        self.backup = {}
        self.backup_eps = {}

    def attack_backward(self, inputs, labels):
        if self.adv_lr == 0:
            return
        self._save()
        self._attack_step()

        y_preds = self.model(inputs)

        adv_loss = self.criterion(y_preds, labels)
        self.optimizer.zero_grad()
        return adv_loss

    def _attack_step(self):
        e = 1e-6
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                norm1 = torch.norm(param.grad)
                norm2 = torch.norm(param.data.detach())
                if norm1 != 0 and not torch.isnan(norm1):
                    # 在损失函数之前获得梯度
                    r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)
                    param.data.add_(r_at)
                    param.data = torch.min(
                        torch.max(param.data, self.backup_eps[name][0]), self.backup_eps[name][1]
                    )

    def _save(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                if name not in self.backup:
                    self.backup[name] = param.data.clone()
                    grad_eps = self.adv_eps * param.abs().detach()
                    self.backup_eps[name] = (
                        self.backup[name] - grad_eps,
                        self.backup[name] + grad_eps,
                    )

    def _restore(self,):
        for name, param in self.model.named_parameters():
            if name in self.backup:
                param.data = self.backup[name]
        self.backup = {}
        self.backup_eps = {}

2.2 经典训练模式

for step, batch in enumerate(train_loader):
    inputs, labels = batch
    
    # 将模型的参数梯度初始化为0
    optimizer.zero_grad()
    
    # forward + backward + optimize
    predicts = model(inputs)          # 前向传播计算预测值
    loss = loss_fn(predicts, labels)  # 计算当前损失
    loss.backward()       # 反向传播计算梯度
	loss.backward()
    optimizer.step()                  # 更新所有参数 

2.3 awp训练模式

# 初始化AWP
awp = AWP(model, loss_fn, optimizer, adv_lr=awp_lr, adv_eps=awp_eps)

for step, batch in enumerate(train_loader):
    inputs, labels = batch
    
    # 将模型的参数梯度初始化为0
    optimizer.zero_grad()
    
    # forward + backward + optimize
    predicts = model(inputs)          # 前向传播计算预测值
    loss = loss_fn(predicts, labels)  # 计算当前损失
    loss.backward()       # 反向传播计算梯度
    # 指定从第几个epoch开启awp，一般先让模型学习到一定程度之后
    if awp_start >= epoch:
        loss = awp.attack_backward(inputs, labels)
        loss.backward()
        awp._restore()                    # 恢复到awp之前的model
    optimizer.step()                  # 更新所有参数 

2.4 超参数

• adv_param (str): 要攻击的layer name，一般攻击第一层或者全部weight参数效果较好
• adv_lr (float): 主要调整的参数，表示攻击步长，该参数相对难调节，如果只攻击第一层embedding，一般用1比较好，全部参数用0.1比较好
• adv_eps (float): 主要调整参数，表示参数扰动最大幅度限制，范围（0 ~ +∞），一般设置（0，1）之间相对合理一点
• start_epoch (int): （0 ~ +∞）什么时候开始扰动，默认是0，如果效果不好可以调节值模型收敛一半的时候再开始攻击

参考文献

• Adversarial Distributional Training for Robust Deep Learning)
• paper code
• https://zhuanlan.zhihu.com/p/563641649