FGM对抗训练

import torch
import numpy as np
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as Data
from torch.utils.data import Dataset, DataLoader

sentence_list = ['i say hello', 'you eat egg', 'he can sing', 'she will go']

words = ' '.join(sentence_list).split(' ')
all_words = list(set(words))
print(all_words)
all_words_len = len(all_words)
print(all_words_len)

# 构建i2v，v2i
index2word = {k: v for k, v in enumerate(all_words)}
word2index = {v: k for k, v in enumerate(all_words)}
print(index2word)
print(word2index)

## 构建x(输入值)和y(预测值)，并把x(输入值)和y(预测值)转为one-hot编码
def make_data(sentence_list):
    input_batch = []
    output_batch = []
    for sentence in sentence_list:
        # 将句子切分为单词,然后再将前两个单词作为x(输入值),最后一个单词作为y(预测值)
        word = sentence.split(' ')
        input = [word2index[n] for n in word[:-1]]
        output = word2index[word[-1]]
        # np.eye(all_words_len)[input]可以理解为,将每个单词的编号转为 one-hot编码
        # 输入值转为one-hot编码,输出值不用转,因为相当于13分类问题(一共有13种不同的单词)
        input_batch.append(np.eye(all_words_len)[input])
        output_batch.append(output)
    print(input_batch)
    print(output_batch)
    return torch.Tensor(input_batch), torch.LongTensor(output_batch)

input_batch, output_batch = make_data(sentence_list)
print('-------')
print(input_batch, output_batch)

## 构建dataset和dataloader
dataset = Data.TensorDataset(input_batch, output_batch)
dataloader = Data.DataLoader(dataset=dataset,
                             batch_size=2,
                             drop_last=True,
                             shuffle=False)
#shuffle是否打乱顺序，默认为False表示不打乱顺序
# drop_last表示是否舍弃最后一个，
#举例：本次总的数据是五个，每次运送两个数据,最后一句话,也就是'i say ok'不参加训练


for x, y in dataloader:
    print(x, x.shape)
    # [2, 2, 13]
    # 第一个2代表2个数据,也就是batch_size,
    # 第二个2表示x(输入值)中单词的个数,本次是由前两个单词,来预测第三个单词
    # 13 表示每个单词被编码为13维
    print(y, y.shape)
    print('----')
    # 可以发现是两个一组进行输出,因为batch_size=2
    ### 此时数据处理部分内容就已经完成,接下来构建RNN网络,将数据用dataloader送入到网络中

####构建RNN网络
class MyRNN(nn.Module):
    def __init__(self):
        super(MyRNN, self).__init__()
        #all_words_len=13,也就是sentence_list中,总共有13种不同的单词,
        # 也就预测最后一个单词,总共有13中不同的结果,也就是13分类问题
        self.rnn = nn.RNN(input_size=all_words_len, hidden_size=6)
        self.fc = nn.Linear(in_features=6, out_features=all_words_len)

    #图片中使out,ht=forward(x,h0)
    # 也就是 out,hidden = self.rnn(x,hidden)
    # 这里采用input代表x，h_i代表h0(hidden)
    # 就会变成out,ht=self.rnn(input,h_0) # rnn中数据放在第一个位置
    def forward(self, input):
        input = input.transpose(0, 1)
        # input是dataloader中的x数据,x数据的格式是
        # RNN的前向传播,需要传入两个参数
        out, h_0i = self.rnn(input)
        out = out[-1]
        model = self.fc(out)
        return model

#实例化对象
model = MyRNN()

class FGM():
    def __init__(self, model):
        self.model = model
        self.backup = {}

    def attack(self, epsilon=1., emb_name='emb'):
        # emb_name这个参数要换成你模型中embedding的参数名
        # 例如，self.emb = nn.Embedding(5000, 100)
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)  # 默认为2范数
                if norm != 0:
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='emb'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

## 定义损失函数和优化器## 固定步骤
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 初始化
fgm = FGM(model)

for epoch in range(10):
    for x, y in dataloader:
        # h_0i = torch.randn(1, x.shape[0], 6)
        # print(h_0i, h_0i.shape)  # torch.Size([1, 2, 6])
        out1= model(x)
        loss=loss_fn(out1,y)
        # print(torch.zeros(2,12))
        loss.backward()# torch.ones(2,12)
        fgm.attack()  # embedding被修改了
        # print('111')
        out1 = model(x)
        loss_sum=loss_fn(out1,y)  #torch.ones(2,12) 反向传播，在正常的grad基础上，累加对抗训练的梯度
        # 梯度下降，更新参数
        # optimizer.step()
        loss_sum.backward()
        fgm.restore()  # 恢复Embedding的参数
        optimizer.step()
        model.zero_grad()
        print(loss)
        print(loss_sum)

correct = 0
total = 0
with torch.no_grad():
    for x, y in dataloader:
        outputs = model(x)
        print(outputs)
        _, predicted = torch.max(outputs.data, 1)
        total += y.size(0)
        correct += (predicted == y).sum().item()
        print(correct)
print('Accuracy of the network : %d %%' % (100 * correct / total))