小白对抗训练入门(1)--FGM

序

先想再做，对比总结！

原理

FGM核心思想是：在训练时对样本施加一定的变形，从而提升模型的健壮性。

如何施加一定的变形呢，而且还不会把模型搞为白痴呢，而且施加的形变太小，容易没啥用，施加的太大又容易搞错边界。作者提出对每个样本施加一个梯度变化方向上的$ϵ$，即可以解决此问题。具体步骤如下：

1. 对于训练样本$x$和$y$, 得出来预测值$f(x)$，并求梯度$\Delta f(x)$。
2. 对于每个特征$x_i$加一个$ϵ\ast sign(\Delta f(x))$，公式如下：$$x^{\prime }=x+ϵ\ast sign(\Delta f(x))$$
3. 对于$x^{\prime }$进行预测得到$f(x^{\prime })$，然后求的梯度，进行反向传播，更新模型参数。

实现

在NLP领域，每个样本最后会变成embedding，输入到模型进行预测。所以只需要对embedding层进行施加变化就好。
具体代码如下：

class PGM:
	def __init__(self, model, epsilon, embedding_name):
		self.model = model
		self.epsilon = epsilon
		self.embedding_name = embedding_name
	
	def attack(self):
		# 找到那一层，然后施加扰动
		for name, para in self.model.named_parameters():
			if name == self.embedding_name and para.requires_grad:
				# 对这个进行embedding
				self.backup[name] = para.data.clone()
				norm = torch.norm(para.grad)
				if norm!=0 and not torch.isnan(norm):
					rat = self.epsilon * (para.grad/norm)
					para.data.add_(rat)
	
	def restore(self):
		for name, para in self.model.named_parameters():
			if name == self.embedding_name:
				para.data = self.backup[name]

# 训练调用
pgm_model = PGM(model, epsilon=1, embedding_name='embedding')

for batch_input in train_dataset:
    # 第一次正向传播 求 f(x)
	pred = model(batch_input['input'])
	# 反向传播，求梯度
	loss = myloss(pred, batch_input['y'])
	loss.backward()
	
	# 开始攻击,得到x'
	pgm_model.attack()
	# 再次方向传播，得到f(x')
	pred_2 = model(batch_input['input'])
	loss2 = myloss(pred_2, batch_input['y'])
	loss2.backward()
	# 回退原始参数
	pgm_model.restore()
	
	# 总体开始更新参数
	optimizer.step()
	model.zero_grad()