Pytorch实现Transformer字符级机器翻译

前言

上次发了一篇用seq2seq with attention做机器翻译的博客，今天我们试试Transformer。这篇文章主要介绍Transformer模型的搭建和训练，数据集仍然是上次的博客中使用的数据集。没看过那篇博客的可以先看看它构建数据集的部分再回来看这篇文章。

搭建模型

重点分析

先看看这张经久不衰的Transformer架构图

实现Transformer有几个重点

Transformer中的三个mask
Multi-Head Attention With Add&Norm
Feed Forward With Add&Norm
Positional Encoding

三个mask

encoder self-attention mask

我们知道为了方便矩阵并行计算，encoder和decoder的输入都是经过PAD的，但是我们不希望在注意力层把PAD考虑进去，所以会通过mask机制去除PAD的影响。具体怎么做呢，举个例子，假如我们有一个句子，经过pad后的长度是5

H i <PAD> <PAD> <PAD>

然后这个句子要经过self-attention计算每个字彼此的相关程度

	H	i
H	1	1
i	1	1
<PAD>	1	1
<PAD>	1	1
<PAD>	1	1

我们横着看这个表，因为在multi-head attention层计算score的时候我们是对最后一维进行softmax。

然后我们可以看到如果没有mask的话，H这个字会受到"Hi"以及后面三个的影响。要消除这个的影响，Transformer使用的方法是在计算softmax之前，让所在位置的值无穷小，从而经过softmax之后就无限接近于0。这样"H"就基本取决于"Hi"了。在实现上即将生成的attention_score与上面的mask做element-wise计算，然后将矩阵中为0的地方替换成-inf。

decoder mask-attention mask

decoder中有一个Masked Multi-Head Attention，这里的Mask的作用就是防止网络偷看答案。比如我们把Hi翻译成你好，那么decoder在翻译第一个字你的时候，就只能看到开始字符，不能看到你。

	<BOS>	你	好	<PAD>	<PAD>	<PAD>
<BOS>	1	0	0	0	0	0
你	1	1	0	0	0	0
好	1	1	1	0	0	0
<PAD>	1	1	1	1	0	0
<PAD>	1	1	1	1	1	0
<PAD>	1	1	1	1	1	1

如上图，通过这样的上三角mask就能实现每个字只受到前面的字的影响。

当然，部分依然还是需要mask掉，这么一来，将两个mask加起来就是最终的mask。

	<BOS>	你	好
<BOS>	1	0	0
你	1	1	0
好	1	1	1
<PAD>	1	1	1
<PAD>	1	1	1
<PAD>	1	1	1

decoder interaction attention mask

我们注意到decoder中除了Masked Multi-Head Atttention之外还有一个Multi-Head Attention，并且这个attention模块除了接收上一个attention的输出之外，还需要接收encoder的输出。这个attention模块的功能就是计算encoder的输出的对decoder 输入的影响。如下表

	H	i
<BOS>	1	1
你	1	1
好	1	1
<PAD>	1	1
<PAD>	1	1
<PAD>	1	1

可以看到这个mask跟encoder self-attention mask类似，只不过这个行列不一样了。

Mask生成代码

针对PAD的mask生成代码如下，同时输入encoder_input和decoder_input。

如果是encoder self-attention mask，那么encoder_input和decoder_input都传为encoder_input。

如果是decoder mask-attention mask，那么encoder_input和decoder_input都传为decoder_input。

如果是decoder interaction attention mask，那么就分别传入encoder_input和decoder_input。

def generate_attention_mask(encoder_input, decoder_input):sd = decoder_input.shape[1]# torch.tile将tensor按照dims中的数字扩充维度# unsqueeze(dim=1) 增加时间序列维度attention_mask = torch.tile(encoder_input.eq(0).unsqueeze(dim=1), dims=(1, sd, 1))return attention_mask

针对decoder self-attention mask中的上三角mask，生成代码如下

def generate_triu_mask(decoder_input):return torch.triu(torch.ones(decoder_input.shape[1], decoder_input.shape[1], device=decoder_input.device), diagonal=1)

然后对decoder mask-attention mask，我们需要把pad mask和triu mask加起来

mask_attention_mask = generate_attention_mask(decoder_input,decoder_input)+generate_triu_mask(decoder_input)

Multi-Head Attention With Add&Norm

实现attention的代码挺多的，在这儿主要讲两点。

一是Multi-Head Atttention不改变数据的维度，即输入时embedding_size，中间处理过程会变成hidden_size,最后输出时依然是embedding_size

二是我们前面生成的mask的维度是b,s,s，但是对于Multi-Head Attention来说生成的attention_score维度是b,n,s,s。n是头的数量。所以通过unsqueeze增加头维度，然后通过torch.tile重复把mask的维度变成b,n,s,s。随后可以直接通过masked_fill_方法进行mask处理

代码如下

class MultiHeadAttention(nn.Module):def __init__(self, embedding_size, hidden_size, num_heads, dropout=0.1):super().__init__()self.num_heads = num_headsself.hidden_size = hidden_sizeself.dk = self.hidden_size // self.num_headsassert self.hidden_size % self.num_heads == 0self.Q = nn.Linear(embedding_size, hidden_size)self.K = nn.Linear(embedding_size, hidden_size)self.V = nn.Linear(embedding_size, hidden_size)self.dropout = nn.Dropout(dropout)self.dense = nn.Linear(hidden_size, embedding_size)self.layernorm = nn.LayerNorm(embedding_size)def forward(self, embedding, encoder_output, attention_mask):residual = embeddingq = self.Q(embedding)k = self.K(encoder_output)v = self.V(encoder_output)b, s, h = q.shapeq = q.view(b, -1, self.num_heads, self.dk).transpose(1, 2)k = k.view(b, -1, self.num_heads, self.dk).transpose(1, 2)v = v.view(b, -1, self.num_heads, self.dk).transpose(1, 2)attention_score = torch.matmul(q, k.transpose(-1, -2))attention_score = attention_score / math.sqrt(self.dk)attention_score = attention_score.masked_fill_(torch.tile(attention_mask.unsqueeze(1), dims=(1, self.num_heads, 1, 1)), float("-inf"))attention_score = torch.softmax(attention_score, dim=-1)attention_score = self.dropout(attention_score)context = torch.matmul(attention_score, v)context = context.transpose(1, 2).contiguous().view(b, s, -1)output = self.dense(context) + residualoutput = self.layernorm(output)return output, attention_score

Feed Forward With Add&Norm

前馈神经网络的实现很简单，在这儿用卷积代替了全连接，从而可以适应任意长度的序列

class PoswiseFeedForward(nn.Module):def __init__(self, in_channels, out_channels):super().__init__()self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=1)self.conv2 = nn.Conv1d(out_channels, in_channels, kernel_size=1)self.layernorm = nn.LayerNorm(in_channels)self.relu = nn.ReLU()def forward(self, x):y = x.transpose(-1, -2)y = self.relu(self.conv1(y))y = self.relu(self.conv2(y))y = y.transpose(-1, -2)y = self.layernorm(y)return x + y

注意在卷积之前和之后需要交换最后两个维度，对于1维卷积通道维度应该在倒数第二维

Positional Encoding

Postional Encoding就基本是按照公式写了，我们只需要明确输入输出的维度变化即可。

输入:input embedding

输出:一个max_len,embedding_dim的张量，max_len维度表示位置，embedding_dim维度即每个位置对应的编码

代码如下

class PositionalEncoding(nn.Module):def __init__(self, dropout=0.1):super().__init__()self.input_size = input_sizeself.dropout = nn.Dropout(dropout)max_len = 5000pe = torch.zeros(size=(max_len, self.input_size))pos = torch.arange(max_len)i2s = torch.arange(self.input_size) // 2 * 2item = pos.unsqueeze(1) * torch.exp(i2s * self.input_size * math.log(10000))pe[:, 0::2] = torch.sin(item)[:, 0::2]pe[:, 1::2] = torch.cos(item)[:, 1::2]pe = pe.unsqueeze(0)self.register_buffer('pe', pe)def forward(self, input_embedding):x = input_embedding + self.pe[:, :input_embedding.shape[1]]x = self.dropout(x)return x

训练

transformer的训练跟seq2seq不同，transformer对超参数非常敏感，如果按照之前seq2seq的训练方法，模型可能很难收敛。所以呢，我参考Bert的训练方式，大家可以看我之前写的Bert模型做多标签文本分类这篇博客。

首先是学习率初始化为2e-4，并且使用warmup和学习率衰减策略。训练200轮，warmup_proportion是0.1

def lr_lambda(step):# 线性变换，返回的是某个数值x，然后返回到类LambdaLR中，最终返回old_lr*xif step < warmup_steps:  # 增大学习率return float(step) / float(max(1, warmup_steps))# 减小学习率return max(0.0, float(t_total - step) / float(max(1.0, t_total - warmup_steps)))

其次是使用AdamW优化器，给除bias和Layernorm.weight的模型参数添加权重衰减，eps初始化为1e-9

param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],'weight_decay': opt.weight_decay},{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=opt.lr, eps=opt.epsilon)

其它如梯度裁剪之类的策略依然沿用。完整代码见章末链接

推理

Transformer的推理过程跟Seq2seq不太一样，因为Tranformer需要输入前面所有生成的字，而seq2seq的decoder推理时只需要输入前一个字即可。

def inference(self, sentence, en_tokenizer, ch_tokenizer, max_length=50):with torch.no_grad():encoder_input = torch.tensor([en_tokenizer.encode(sentence)], device=self.device)encoder_output, attention_states = self.encoder(encoder_input)decoder_input = torch.tensor([[ch_tokenizer.BOS]], device=self.device)result = ""for i in range(max_length):decoder_output = self.decoder(decoder_input, encoder_input, encoder_output)logits = self.classifier(decoder_output)windex = logits.argmax(dim=-1)if windex[0][-1] == 2:breakdecoder_input = torch.cat((decoder_input, windex[:, -1].unsqueeze(0)), dim=1)result += ch_tokenizer.decode(windex[0][-1])return result

全部代码我放在csdn了，见链接