Press "Enter" to skip to content

通过极简翻译模型Demo,彻底理解Transformer

前言:Transformer并没有特别复杂,但是理解transformer对于初学者不是件容易的事,根因在于:Transformer的解读往往没有配套的demo,而缺少端到端的demo,就很难透彻理解Transformer的具体运算流程。github上使用Transformer的翻译模型、推断模型,作为demo来说,代码又太过复杂,不易上手。

 

本文共三个部分:

 

 

    1. 子模块解读 :拆解Transformer,结合代码解读各个子模块的运算细节;

 

    1. 翻译模型Demo :讲解使用transformer的翻译模型,将 (‘<bos>’, ‘i’, ‘am’, ‘iron’, ‘man’, ‘<eos>’) 翻译为 (‘<bos>’, ‘我’, ‘是’, ‘钢铁’, ‘侠’, ‘<eos>’) 的训练与推理过程 。 (训练与推理,都只翻译这一句话);

 

    1. 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。

 

求星星~

 

一、子模块解读

 

 

    1. 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,计算流分别为:

 

 

    1. MultiheadAttention,此处Query、Key、Value都是同一Input,所以也是Self-Attention,后接LayerNorm;

 

    1. 一个shape为 emb_dim, dim_feedforward 的全连接层;

 

    1. 一个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运算流程:

 

 

    1. tgt进行self-attention,经过LayerNorm,当作MultiheadAttention的Query;

 

    1. memory当作MultiheadAttention的Key、Value,结合Query,进行一次MultiheadAttention,经过LayerNorm后输出;

 

    1. 经过一个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。

 

不着急看模型,先看看数据:

 

 

    1. 明确翻译任务

 

 

假设我们已经做好了英文词典和中文词典,并且对每一个字符编号:

 

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的步骤为:

 

 

    1. 将tgt和src经过word_emb得到各个词的词向量,并且经过PositionalEncoding,得到含有position信息的src tensor和tgt tensor;

 

    1. 将src经过TransformerEncoder得到memory;

 

    1. 将tgt与memory送入TransformerDecoder得到tgt_out;

 

    1. 最后经过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矩阵

    1. 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

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注