前言:Transformer并没有特别复杂,但是理解transformer对于初学者不是件容易的事,根因在于:Transformer的解读往往没有配套的demo,而缺少端到端的demo,就很难透彻理解Transformer的具体运算流程。github上使用Transformer的翻译模型、推断模型,作为demo来说,代码又太过复杂,不易上手。
本文共三个部分:
- 子模块解读 :拆解Transformer,结合代码解读各个子模块的运算细节;
- 翻译模型Demo :讲解使用transformer的翻译模型,将 (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’) 翻译为 (‘<bos>’, ‘我’, ‘是’, ‘钢铁’, ‘侠’, ‘<eos>’) 的训练与推理过程 。 (训练与推理,都只翻译这一句话);
- Attention的mask作用: 解读attention中mask的作用。
本文配套的翻译模型Demo,源码地址(本文中的配图高清版,也全部包含在了img文件夹中):
thisiszhou/Transformer-Translate-Demo github.com
Demo中Transformer的基础模块(MultiheadAttention,TransformerEncoder,TransformerDecoder),全部摘自torch.nn.Transformer,为了容易上手,删减了很多参数和分支,很大程度降低了阅读源码的难度,强烈建议下载源码,配套本文阅读。Demo中的Transformer仅供学习原理,工业化使用Transformer,建议使用torch.nn.Transformer。
求星星~
一、子模块解读
- MultiheadAttention(Demo中model.multi_head_attention.MultiheadAttention)
MultiheadAttention多头注意力,和注意力Attention稍有区别,是整个Transformer的核心,其他模块都是MultiheadAttention的封装与组合。下面先讲解Attention,暂时忽略batch_size,可以暂且理解其为batch_size等于1的特例。
Attention计算流程,上图:
图一 Attention
输入为Q、K、V三个矩阵,其中 tgt 指target,src 指source,代表的含义在不同任务中有所差异。在本文的demo英译汉这个任务中,英语 (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’) 就是src,汉语 (‘<bos>’, ‘我’, ‘是’, ‘钢铁’, ‘侠’, ‘<eos>’) 就是tgt。其中,tgt_size,src_size分别指汉语句子的最大长度以及英文句子的最大长度,emb_dim是每个单词词向量的维度。Attention模型进行了如下操作:
(1) 将Q、K、V经过一层全连接层,得到新的Q‘、K‘、V‘;
(2) Q‘与K‘的转置进行矩阵乘,得到矩阵W,将W进行dim=-1维度的softmax操作,得到权重矩阵W‘,矩阵的每一行i,代表了tgt中第i个词,对src每个词的注意力权重,所以W‘的维度为 tgt_size * src_size;
(3) 使用W‘ 矩阵乘V‘,得到新的矩阵O,将O经过一层全连接层,得到输出Output。
以上为Attention的计算流程,Self-Attention,就是Q、K、V输入为同一矩阵,即可计算矩阵关于自己的Attention。
何为 MultiheadAttention?很多解读说是并行多个Attention,之后再合并起来,这种说法不完全对,MultiheadAttention是对emb_dim维度进行切分,之后并行Attention,精髓就在于对emb_dim维度进行切分。
MultiheadAttention计算流程,上图:
图二 MultiheadAttention
同样这里忽略batch_size维度,计算步骤如下:
(1) 将Q、K、V经过一层全连接层,得到新的Q‘、K‘、V‘;
(2)将Q‘、K‘、V‘从emb_dim维度进行切分,共切割num_heads个,这里要求emb_dim可以被num_heads整除;切分后的结果为Q‘1、Q‘2、… 、Q‘n,K‘1、K‘2、… 、K‘n,V‘1、V‘2、… 、V‘n;
(3) Q‘1与K‘1的转置进行矩阵乘,得到矩阵W1,Q‘2与K’2的转置进行矩阵乘,得到矩阵W2,… 。再将每一个W1、W2,… ,进行dim=-1维度的softmax操作得到W‘1、W’2,… ;
(4) 使用W‘1 矩阵乘V‘1,得到矩阵O1,使用W‘2 矩阵乘V‘2,得到矩阵O2,…;
(5)Oi的shape为tgt_size * head_size (这里 ),一共有num_heads个Oi,将这些Oi从head_size维度拼接起来,得到O,shape为tgt_size * emb_dim;
(6)O经过一层全连接层,得到输出Output。
废话不多说,过一遍pytorch版的multiheadattention代码(见demo中的model/multi_head_attention.py)
注释写在代码上方,先忽略mask相应操作,key_padding_mask与attn_mask都作为None处理即可,后文会对mask进行单独讲解。
__init__:一些参数、神经网络层的初始化
class MultiheadAttention(Module): def __init__(self, word_emb_dim, nheads, dropout_prob=0. ): super(MultiheadAttention, self).__init__() self.word_emb_dim = word_emb_dim self.num_heads = nheads self.dropout_prob = dropout_prob self.head_dim = word_emb_dim // nheads assert self.head_dim * nheads == self.word_emb_dim # embed_dim must be divisible by num_heads self.q_in_proj = Linear(word_emb_dim, word_emb_dim) self.k_in_proj = Linear(word_emb_dim, word_emb_dim) self.v_in_proj = Linear(word_emb_dim, word_emb_dim) self.out_proj = Linear(word_emb_dim, word_emb_dim)
MultiheadAttention的前向过程(若代码格式阅读体验不好,可直接查阅完整代码):
def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None): """ :param query: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim] :param key: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] :param value: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] :param key_padding_mask: Tensor, shape: [batch_size, src_sequence_size] :param attn_mask: Tensor, shape: [tgt_sequence_size, src_sequence_size] :return: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim] """ # 获取query的shape,这里按照torch源码要求,按照tgt_sequence_size, batch_size, word_emb_dim顺序排列 tgt_len, batch_size, word_emb_dim = query.size() num_heads = self.num_heads assert word_emb_dim == self.word_emb_dim head_dim = word_emb_dim // num_heads # 检查word_emb_dim是否可以被num_heads整除 assert head_dim * num_heads == word_emb_dim scaling = float(head_dim) ** -0.5 # 三个Q、K、V的全连接层 q = self.q_in_proj(query) k = self.k_in_proj(key) v = self.v_in_proj(value) # 这里对Q进行一个统一常数放缩 q = q * scaling # multihead运算技巧,将word_emb_dim切分为num_heads个head_dim,并且让num_heads与batch_size暂时使用同一维度 # 切分word_emb_dim后将batch_size * num_heads转换至第0维,为三维矩阵的矩阵乘法(bmm)做准备 q = q.contiguous().view(tgt_len, batch_size * num_heads, head_dim).transpose(0, 1) k = k.contiguous().view(-1, batch_size * num_heads, head_dim).transpose(0, 1) v = v.contiguous().view(-1, batch_size * num_heads, head_dim).transpose(0, 1) src_len = k.size(1) # Q、K进行bmm批次矩阵乘法,得到权重矩阵 attn_output_weights = torch.bmm(q, k.transpose(1, 2)) assert list(attn_output_weights.size()) == [batch_size * num_heads, tgt_len, src_len] if attn_mask is not None: if attn_mask.dtype == torch.bool: attn_output_weights.masked_fill_(attn_mask, float('-inf')) else: attn_output_weights += attn_mask if key_padding_mask is not None: attn_output_weights = attn_output_weights.view(batch_size, num_heads, tgt_len, src_len) attn_output_weights = attn_output_weights.masked_fill( key_padding_mask.unsqueeze(1).unsqueeze(2), float('-inf'), ) attn_output_weights = attn_output_weights.view(batch_size * num_heads, tgt_len, src_len) # 权重矩阵进行softmax,使得单行的权重和为1 attn_output_weights = torch.softmax(attn_output_weights, dim=-1) # print(f"attn_output_weights: {attn_output_weights}") attn_output_weights = torch.dropout(attn_output_weights, p=self.dropout_prob, train=self.training) # 权重矩阵与V矩阵进行bmm操作,得到输出 attn_output = torch.bmm(attn_output_weights, v) assert list(attn_output.size()) == [batch_size * num_heads, tgt_len, head_dim] # 转换维度,将num_heads * head_dim reshape回word_emb_dim,并且将batch_size调回至第1维 attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, batch_size, word_emb_dim) # 最后一层全连接层,得到最终输出 attn_output = self.out_proj(attn_output) return attn_output
2. TransformerEncoder(Demo中model.transformer_encoder)
Transformer主要有TransformerEncoder和TransformerDecoder组成,一个是编码,一个解码。编码时只对src进行编码,例如本文中的例子,将英文翻译为中文,那幺只对src英文输入语句进行encode。
TransformerEncoder计算流程,上图(图中省略激活层Relu以及Dropout):
图三 TransformerEncoder
TransformerEncoder由6个(数量可自定义)TransformerEncoderLayer组成,一个TransformerEncoderLayer,计算流分别为:
- MultiheadAttention,此处Query、Key、Value都是同一Input,所以也是Self-Attention,后接LayerNorm;
- 一个shape为 emb_dim, dim_feedforward 的全连接层;
- 一个shape为dim_feedforward, emb_dim的全连接层,后接LayerNorm,输出Output作为当前TransformerEncoderLayer的输出;
由于TransformerEncoderLayer的输入shape为src_size * emb_dim,输出shape也为src_size * emb_dim,所以TransformerEncoderLayer的输出可以直接喂给下一个TransformerEncoderLayer。重复六次之后,就得到了TransformerEncoder的输出。
TransformerEncoderLayer的代码,相对简单:
__init__:
class TransformerEncoderLayer(Module): def __init__(self, word_emb_dim, nhead, dim_feedforward=2048, dropout_prob=0.1): super(TransformerEncoderLayer, self).__init__() self.self_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob) self.linear1 = Linear(word_emb_dim, dim_feedforward) self.dropout = Dropout(dropout_prob) self.linear2 = Linear(dim_feedforward, word_emb_dim) self.norm1 = LayerNorm(word_emb_dim) self.norm2 = LayerNorm(word_emb_dim) self.dropout1 = Dropout(dropout_prob) self.dropout2 = Dropout(dropout_prob) self.activation = torch.relu
前向过程:
def forward(self, src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: """ :param src: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] :param src_mask: Tensor, shape: [src_sequence_size, src_sequence_size] :param src_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size] :return: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] """ # self attention src2 = self.self_attn(src, src, src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask) src = src + self.dropout1(src2) src = self.norm1(src) # 两层全连接 src2 = self.linear2(self.dropout(self.activation(self.linear1(src)))) src = src + self.dropout2(src2) src = self.norm2(src) return src
而TransformerEncoder就是6个TransformerEncoderLayer的copy:
__init__:
class TransformerEncoder(Module): __constants__ = ['norm'] def __init__(self, encoder_layer, num_layers, norm): super(TransformerEncoder, self).__init__() # 将同一个encoder_layer进行deepcopy n次 self.layers = _get_clones(encoder_layer, num_layers) self.num_layers = num_layers self.norm = norm
前向过程:
def forward(self, src: Tensor, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None) -> Tensor: """ :param src: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] :param mask: Tensor, shape: [src_sequence_size, src_sequence_size] :param src_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size] :return: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] """ output = src # 串行n个encoder_layer for mod in self.layers: output = mod(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask) output = self.norm(output) return output
3. TransformerDecoder(Demo中model.transformer_decoder)
先上图:
图四 TransformerEncoder
Decoder有两个输入,一个是Encoder的输出memory,一个是tgt。这里tgt读者可能会有问题,Decoder对Encoder的输出memory进行解码理所应当,但是tgt是哪来的?首先,没有tgt输入,只对memory解码的Decoder也是存在的,但是在翻译任务中,需要有一个额外的tgt输入,来得到不同的输出。当前先讲解Decoder的运算流程,这里tgt具体应用,在第二章翻译demo中会详细解释。
TransformerDecoderLayer运算流程:
- tgt进行self-attention,经过LayerNorm,当作MultiheadAttention的Query;
- memory当作MultiheadAttention的Key、Value,结合Query,进行一次MultiheadAttention,经过LayerNorm后输出;
- 经过一个shape为 emb_dim, dim_feedforward 的全连接层,和一个shape为dim_feedforward, emb_dim的全连接层,后接LayerNorm,输出tgt_out,shape与最开始的输入tgt相同。
TransformerDecoder就是六个(数量可自定义)TransformerDecoderLayer串联组成。所以重点看TransformerDecoderLayer实现,上代码:
__init__:
class TransformerDecoderLayer(Module): def __init__(self, word_emb_dim, nhead, dim_feedforward=2048, dropout_prob=0.1): super(TransformerDecoderLayer, self).__init__() # 初始化基本层 self.self_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob) self.multihead_attn = MultiheadAttention(word_emb_dim, nhead, dropout_prob=dropout_prob) # Implementation of Feedforward model self.linear1 = Linear(word_emb_dim, dim_feedforward) self.dropout = Dropout(dropout_prob) self.linear2 = Linear(dim_feedforward, word_emb_dim) self.norm1 = LayerNorm(word_emb_dim) self.norm2 = LayerNorm(word_emb_dim) self.norm3 = LayerNorm(word_emb_dim) self.dropout1 = Dropout(dropout_prob) self.dropout2 = Dropout(dropout_prob) self.dropout3 = Dropout(dropout_prob) self.activation = torch.relu
前向过程:
def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None) -> Tensor: """ :param tgt: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim] :param memory: Tensor, shape: [src_sequence_size, batch_size, word_emb_dim] :param tgt_mask: Tensor, shape: [tgt_sequence_size, tgt_sequence_size] :param memory_mask: Tensor, shape: [src_sequence_size, src_sequence_size] :param tgt_key_padding_mask: Tensor, shape: [batch_size, tgt_sequence_size] :param memory_key_padding_mask: Tensor, shape: [batch_size, src_sequence_size] :return: Tensor, shape: [tgt_sequence_size, batch_size, word_emb_dim] """ # tgt的self attention tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask) tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt) # tgt与memory的attention tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask) tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt) # 两层全连接层 tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt
4. PositionalEncoding
PositionalEncoding虽然在transformer中是第一个模块,但是这里最后讲解,因为PositionalEncoding对于transformer来说,不是必须的,在torch.nn.Transformer中,没有包含PositionalEncoding,需要自己实现。有一些序列化含义不强的场景,PositionalEncoding可以省略。
在翻译任务中,当英文句子被表示成一个矩阵(每一行是句子中对应位置英文单词的词向量),位置信息被淡化,所以PositionalEncoding的作用,就是体现每个词的相对位置与绝对位置信息。
论文中位置编码公式如下:
注意:PositionalEncoding并不是对句子的位置进行一维编码,而是对句子的位置position、每个单词的词向量emb_dim,共两个维度进行编码,而且是独立编码,两个维度互不干涉。上述公式中,pos就是当前词在句子的位置,2i与2i+1,是词向量emb_dim中的位置。所以PositionalEncoding编码矩阵的shape是 (tgt_sequence_size, word_emb_dim)
直接上代码:
__init__:
class PositionalEncoding(nn.Module): def __init__(self, word_emb_dim: int, dropout=0.1, max_len=5000): super(PositionalEncoding, self).__init__() self.dropout = nn.Dropout(p=dropout) position_emb = torch.zeros(max_len, word_emb_dim) # position 编码 position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) dim_div_term = torch.exp(torch.arange(0, word_emb_dim, 2).float() * (-math.log(10000.0) / word_emb_dim)) # word_emb_dim 编码 position_emb[:, 0::2] = torch.sin(position * dim_div_term) position_emb[:, 1::2] = torch.cos(position * dim_div_term) pe = position_emb.unsqueeze(0).transpose(0, 1) # shape: (max_len, 1, d_model) self.register_buffer('pe', pe)
前向过程:
def forward(self, x: Tensor): """ :param x: Tensor, shape: [batch_size, sequence_length, word_emb_dim] :return: Tensor, shape: [batch_size, sequence_length, word_emb_dim] """ # 编码信息与原始信息加和后输出 x = x + self.pe[:x.size(0), :] return self.dropout(x)
PositionalEncoding位置编码,与原始信息加和后输出。
二、翻译模型Demo
首先声明,Transformer的结构并不是固定的,上面所述的四种基础结构(MultiheadAttention,TransformerEncoder,TransformerDecoder,PositionalEncoding)基本是固定的,而Transformer可以由自由组合,torch.nn.Transformer官方demo中,Decoder就是一个全连接层,而没有用上述TransformerDecoder。
不着急看模型,先看看数据:
- 明确翻译任务
假设我们已经做好了英文词典和中文词典,并且对每一个字符编号:
cn_dict = { '<bos>': 0, '<eos>': 1, '<pad>': 2, '是': 3, '千': 4, '你': 5, '万': 6, '在': 7, '我': 8, '人': 9, '三': 10, '一': 11, '侠': 12, '遍': 13, '二': 14, '爱': 15, '好': 16, '钢铁': 17 } en_dict = { '<bos>': 0, '<eos>': 1, '<pad>': 2, 'i': 3, 'three': 4, 'am': 5, 'love': 6, 'you': 7, 'he': 8, 'times': 9, 'is': 10, 'thousand': 11, 'hello': 12, 'iron': 13, 'man': 14 }
语料也已经准备好了,只有如下两句话,每个单词都被收录在上述词典中:
sentence_pair_demo = [ [ ('<bos>', 'i', 'am', 'iron', 'man', '<eos>'), ('<bos>', '我', '是', '钢铁', '侠', '<eos>') ], [ ('<bos>', 'i', 'love', 'you', 'three', 'thousand', 'times', '<eos>'), ('<bos>', '我', '爱', '你', '三', '千', '遍', '<eos>') ] ]
本文翻译模型demo,只训练翻译两句话。了解Transformer最简易翻译模型的原理,只包含两句话的数据集足够。真正使用时,只需要替换更大的训练数据集即可。
注:当前例句假设已经分好词,并且一句话有完整的开始标记符'<bos>’,和结束标记符'<eos>’。
2. 基于Transformer的翻译模型结构
图五 Transformer翻译模型
transformer的步骤为:
- 将tgt和src经过word_emb得到各个词的词向量,并且经过PositionalEncoding,得到含有position信息的src tensor和tgt tensor;
- 将src经过TransformerEncoder得到memory;
- 将tgt与memory送入TransformerDecoder得到tgt_out;
- 最后经过reshape与一层全连接层,得到一个长度为tgt词典长度的向量,代表每个词此时的预测概率(需要经过softmax)
注意:src与tgt通过<pad>填充为固定长度,模型每次输入翻译前原句,和翻译后已经获得的词,来预测下一个词,例如:输入为 src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘我’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),那幺模型下一个预测的词,应该为 ‘是’。
3. 训练(详见no_mask_train.py)
就算是只有两句话的数据集,也需要有一个Dataset工具来组织batch数据(详见代码utils.dataset)。这里Dataset的实现不展开,和Transformer关系不大,只需要明确一下get_batch函数的输出:
def get_batch(self, batch_size=2, padding_str='<pad>', need_padding_mask=False): """ :return: src, tgt_in, tgt_out, src_padding_mask, tgt_padding_mask src: Tensor, shape: [batch_size, src_sequence_size] tgt_in: Tensor, shape: [batch_size, tgt_sequence_size] tgt_out: Tensor, shape: [batch_size, tgt_sequence_size] src_padding_mask: Tensor, shape: [batch_size, src_sequence_size] tgt_padding_mask: Tensor, shape: [batch_size, tgt_sequence_size] """
同样,这里暂时忽略mask。在训练的时候,[(‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’), (‘<bos>’, ‘我’, ‘是’, ‘钢铁’, ‘侠’, ‘<eos>’)],是一组数据,单batch的一些数据示例如下:
(1) src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘我’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),tgt_out = ‘是’;
(2) src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘我’, ‘是’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),tgt_out = ‘钢铁’;
(3)src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘我’, ‘是’, ‘钢铁’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),tgt_out = ‘侠’
如上所示,当前翻译Demo设计的是,输入英文整句,以及中文已经翻译出的词,来预测下一个词。
例如('<pad>’为填充字符,旨在将每句话的长度拉齐):
第一步 :输入为src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),期望输出为 ‘我’;
第二步 :输入为src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘我’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),期望输出为’是’;
直到整句话输出为'<eos>’或者达到最大长度后停止。
那幺在transformer的训练中,输入是src = (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’, ‘<pad>’, ‘<pad>’),tgt = (‘<bos>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’, ‘<pad>’),输出的标签就是’我’。当前这里所有的词,都会转换成该词在词典中的序号,并且按照batch拼成张量。一句话的长度是src_sequence_size或者tgt_sequence_size,按照batch输出后,Tensor的shape为[batch_size, src_sequence_size] or [batch_size, tgt_sequence_size]。
4. 预测(详见no_mask_infer.py)
预测最主要的函数如下:
def infer_with_transformer(model: Transformer, src: Tensor, tgt_dict: Dictionary, max_length=8) -> List[str]: out_seq = ['<bos>'] predict_word = '' while len(out_seq) < max_length and predict_word != '<eos>': tgt_in = transform_words_to_tensor(out_seq, tgt_dict) output = model(src, tgt_in) word_i = torch.argmax(output, -1).item() predict_word = tgt_dict.i2w[word_i] out_seq.append(predict_word) return out_seq
在最开始,out_seq只包含一个'<bos>’,每次预测出的词,都append进out_seq,再循环预测下一个词,直到输出为'<eos>’或者达到最大长度后停止。
使用训练好的模型,可以看到以下输出:
Input sentence: ['<bos>', 'i', 'love', 'you', 'three', 'thousand', 'times', '<eos>'] After translate: ['<bos>', '我', '爱', '你', '三', '千', '遍', '<eos>'] Input sentence: ['<bos>', 'i', 'am', 'iron', 'man', '<eos>'] After translate: ['<bos>', '我', '是', '钢铁', '侠', '<eos>']
三、Attention的mask作用
在MultiheadAttention中,一共有两个mask,一个是key_padding_mask,一个是attn_mask。两个mask相互是独立的,和multihead没有太大关系,所以用Attention来讲解mask的作用。先回顾一下Attention:
图六 Attention
Query和Key的作用,是计算每一个tgt关于每一个src的权重,所以W矩阵(注意是在softmax前的W矩阵)的shape为tgt_size * sec_size,两个mask的作用就是体现在这里。
假设当前W矩阵如下(这里用1填充,方便演示计算):
图七 W矩阵
- key_padding_mask
注意到W中,src有填充符号<pad>,该字符并不需要被tgt注意到,因为对于翻译没有作用。而W矩阵中,每一个tgt字符,都有关于<pad>字符的注意力权重,这里只需要一个key_padding_mask矩阵,进行如下操作,即可消除tgt对于src中<pad>的注意力权重:
图八 key_padding_mask计算流程
在MultiheadAttention的forward中,key_padding_mask参数传入的是一个Tensor,shape为[batch_size, src_sequence_size],可以为bool型(padding的位置为True),也可以为float型(padding位置为-inf,其余为0)。
实验(详见key_padding_mask_train.py与key_padding_mask_infer.py):
Training时,在最外层的Transformer model,传入src_key_padding_mask,与tgt_key_padding_mask:
src, tgt_in, tgt_out, src_pad_mask, tgt_pad_mask = dataset.get_batch(batch_size=1, need_padding_mask=True) output = model(src, tgt_in, src_key_padding_mask=src_pad_mask, tgt_key_padding_mask=tgt_pad_mask)
运行key_padding_mask_train.py,可见收敛速度与训练速度都有提升
step: 100 loss: 1.8999741017573326 step: 200 loss: 1.0202430053596518 step: 300 loss: 0.6763626998902951 step: 400 loss: 0.39750630465763587 step: 500 loss: 0.22424202706432 step: 600 loss: 0.140881586470018 step: 700 loss: 0.11815606149311553 step: 800 loss: 0.0592925045802669 step: 900 loss: 0.04602800890285985 step: 1000 loss: 0.03302763575842654 finished train ! model saved file: weights/padding_mask_model.pkl
2. attn_mask
attn_mask与key_padding_mask相互独立,其shape为 [tgt_sequence_size, src_sequence_size],旨在控制tgt关于src的注意力权重。例如,比较常见的是上三角attn_mask矩阵:
图九 attn_mask计算流程
加入上三角attn_mask矩阵后,tgt的第一个词,只能关注到src的第一个词,tgt的第二个词,只能关注到src的第一个词和第二个词,… 。
本文中的demo,并不需要加入attn_mask,因为在预测第二个词时,神经网络的输入只有之前的词。而在有些任务中,前序列元素不需要关注后序列元素,就可以使用上三角attn_mask来控制。
实验(详见key_padding_mask_train.py与key_padding_mask_infer.py):
假设,我们在本文的翻译任务中强行加入一个限制,tgt中的每个词,只能关注tgt自身当前词之前的词,不能关注自身之后的词,可以先获取上三角attn_mask矩阵,同样,float型和bool型Tensor都可以。
获取上三角矩阵:
def get_upper_triangular(size): return torch.triu(torch.ones(size, size), diagonal=1) == 1
在Transformer的训练和推理都加入该上三角矩阵:
output = model(src, tgt_in, tgt_mask=upper_tri)
运行triu_mask_train.py:
step: 100 loss: 2.169212473629063 step: 200 loss: 1.5380522502928162 step: 300 loss: 1.1274102221317484 step: 400 loss: 0.7371829162703203 step: 500 loss: 0.5449878707849697 step: 600 loss: 0.38455579467139256 step: 700 loss: 0.26030085991229446 step: 800 loss: 0.2031404335089682 step: 900 loss: 0.15147940102617086 step: 1000 loss: 0.1201247349717585 finished train ! model saved file: weights/triu_mask_model.pkl
可见,与不添加上三角矩阵的no_mask_train训练差异不大。
再次强烈建议读一下本文提供的demo源码:
thisiszhou/Transformer-Translate-Demo github.com
求赞求星星。
文章有些长,难免有笔误,欢迎指正。
完。
Be First to Comment