Attention is All You Need（Pytorch实现）

最新推荐文章于 2023-04-23 17:00:58 发布

Nlper_Zhao

最新推荐文章于 2023-04-23 17:00:58 发布

阅读量3.2k

点赞数 6

分类专栏：自然语言处理文章标签：自然语言处理 pytorch 深度学习神经网络

原文链接：https://github.com/bentrevett/pytorch-seq2seq/blob/master/6%20-%20Attention%20is%20All%20You%20Need.ipynb

版权

自然语言处理专栏收录该内容

33 篇文章 31 订阅

订阅专栏

文章目录

在本笔记中，我们将实现一个(稍作修改的版本)的 Attention is All You Need 论文中的transformer模型。本笔记本中的所有图像都将取自transformer 论文。有关Transformer的更多信息，请参见这三篇文章。

引言

与卷积序列到序列模型类似，Transformer不使用任何递归，它也不使用任何卷积层。相反，该模型完全由线性层、注意力机制和规范化组成。

截止到2020年1月，transformer是NLP的主要架构和用于实现许多任务的最先进的结果，并且似乎在不久的将来也会如此。

最流行的Transformer变体是BERT(Bidirectional Encoder Representations from Transformers 来自变压器的双向编码器表示)，在NLP模型中，预训练版本的BERT通常用于替换嵌入层(如果不是更多的话)。

处理预先训练过的transformer的一个常用库是transformer库，请参阅这里的所有预先训练过的可用模型列表。

本笔记本的实现与论文的不同之处在于:

我们使用学习到的位置编码不是静态的
我们使用具有静态学习速率的标准Adam优化器，而不是带有热身和冷却步骤的优化器
我们不使用标签平滑

我们做了所有这些改变，因为他们紧密地跟随BERT’s设置和大多数Transformer变体使用类似的设置。

数据预处理

一如既往，让我们导入所有必需的模块，并为可再现性设置随机种子。

import torch
import torch.nn as nn
import torch.optim as optim

import torchtext
from torchtext.datasets import Multi30k
from torchtext.data import Field, BucketIterator

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker

import spacy
import numpy as np

import random
import math
import time

SEED = 1234

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True

然后，我们将像之前一样创建标记器。

spacy_de = spacy.load('de_core_news_sm')
spacy_en = spacy.load('en_core_web_sm')


def tokenize_de(text):
    """
    Tokenizes German text from a string into a list of strings
    """
    return [tok.text for tok in spacy_de.tokenizer(text)]

def tokenize_en(text):
    """
    Tokenizes English text from a string into a list of strings
    """
    return [tok.text for tok in spacy_en.tokenizer(text)]

我们的字段与之前的笔记相同。模型希望先用批处理维度输入数据，所以我们使用batch_first = True。


SRC = Field(tokenize = tokenize_de, 
            init_token = '<sos>', 
            eos_token = '<eos>', 
            lower = True, 
            batch_first = True)

TRG = Field(tokenize = tokenize_en, 
            init_token = '<sos>', 
            eos_token = '<eos>', 
            lower = True, 
            batch_first = True)

然后我们加载Multi30k数据集并构建词汇表。

train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'), 
                                                    fields = (SRC, TRG))

SRC.build_vocab(train_data, min_freq = 2)
TRG.build_vocab(train_data, min_freq = 2)

最后，我们定义设备和数据迭代器。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

BATCH_SIZE = 128

train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
    (train_data, valid_data, test_data), 
     batch_size = BATCH_SIZE,
     device = device)

搭建模型

接下来，我们将构建模型。像以前的笔记一样，它由一个编码器和一个解码器组成，编码器将输入/源句子(德语)编码成上下文向量，解码器对上下文向量进行解码，输出我们的输出/目标句子(英语)。

Encoder

与ConvSeq2Seq模型类似，Transformer的编码器不试图压缩整个源语句， $X = (x_1，…，x_n)$ ，转换为单个上下文向量 $z$ 。相反，它产生一个上下文向量序列，Z = (z_1，…z_n)。因此，如果我们的输入序列有5个标记那么 $Z = (z_1, z_2, z_3, z_4, z_5)$ 。为什么我们称它为上下文向量序列而不是隐藏状态序列?RNN中时间 $t$ 的隐藏状态只看到t标记 $x_t$ 和它之前的所有标记。但是，这里的每个上下文向量都看到了输入序列中所有位置上的所有标记。
在这里插入图片描述
首先，标记通过一个标准嵌入层。其次，由于模型没有循环，它不知道序列中标记的顺序。我们通过使用称为位置嵌入层的第二个嵌入层来解决这个问题。首先是一个标准的嵌入层，其中输入的不是标记本身，而是标记在序列中的位置，从第一个位置为0的标记《sos》(序列开始)标记开始。位置嵌入的“词汇”大小为100，这意味着我们的模型可以接受长达100个符号的句子。如果我们想处理更长的句子，这个比例可以增加。

从原来的Transformer 实现中可以注意到 Attention is All You Need论文上不学习位置的嵌入。相反，它使用一个固定的静态嵌入。现在的Transformer 架构，如BERT，使用位置嵌入代替，因此我们决定在这些教程中使用它们。查看这一节，了解原始Transformer模型中使用的位置嵌入的更多信息。

接下来，将标记和位置的嵌入元素相加，得到一个包含有关标记及其在序列中的位置的信息的向量。然而，在对其求和之前，标记嵌入要先乘以一个比例因子 $\sqrt{d_{model}}$ ，其中 $d_{model}$ 是隐藏维度大小hid_dim。这被认为减少了嵌入的方差，没有这个比例因子，模型很难可靠地训练。然后将Dropout应用于组合嵌入。

组合嵌入然后通过 $N$ 层编码器层得到 $Z$ ，然后输出，它可以被解码器使用。

源掩码src_mask与源句子形状相同，但当源句中的标记不是《pad》标记时，值为1，当它是《pad》标记时值为0。这在编码器层中被用来掩盖多头注意力机制，这些机制被用来计算和应用对源句子的注意力，所以模型不注意《pad》标记，它不包含有用的信息。

class Encoder(nn.Module):
    def __init__(self, input_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, device, max_length=100):
        super().__init__()

        self.device = device

        self.tok_embedding = nn.Embedding(input_dim, hid_dim)
        self.pos_embedding = nn.Embedding(max_length, hid_dim)

        self.layers = nn.ModuleList([EncoderLayer(hid_dim,
                                                  n_heads,
                                                  pf_dim,
                                                  dropout,
                                                  device)
                                     for _ in range(n_layers)])

        self.dropout = nn.Dropout(dropout)

        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)

    def forward(self, src, src_mask):
        # src = [batch size, src len]
        # src_mask = [batch size, 1, 1, src len]

        batch_size = src.shape[0]
        src_len = src.shape[1]

        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)

        # pos = [batch size, src len]

        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos))

        # src = [batch size, src len, hid dim]

        for layer in self.layers:
            src = layer(src, src_mask)

        # src = [batch size, src len, hid dim]

        return src

Encoder Layer

编码器层包含编码器的所有“肉”。我们首先将源句子及其掩码传递到多头注意力层，然后对其进行dropout，再应用一个残差连接，并将其传递到层归一化层。然后，我们通过一个位置上的前馈层，然后，再一次应用dropout，残差连接，然后层归一化，以得到这一层的输出，并馈入下一层。参数在层之间不共享。

多头注意力层是编码器层用来关注源句子的，也就是说，它在计算和应用注意力到自己本身而不是另一个序列上，因此我们称它为自我注意。

这篇文章对层归一化进行了更详细的介绍，但其要点是对特征的值进行归一化，即跨隐维，因此每个特征的均值为0，标准差为1。这使得具有更大层次的神经网络，如Transformer，更容易训练。

class EncoderLayer(nn.Module):
    def __init__(self, hid_dim, n_heads, pf_dim, dropout, device):
        super().__init__()

        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)
        self.ff_layer_norm = nn.LayerNorm(hid_dim)
        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, pf_dim, dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, src_mask):
        # src = [batch size, src len, hid dim]
        # src_mask = [batch size, 1, 1, src len] 

        # self attention
        _src, _ = self.self_attention(src, src, src, src_mask)

        # dropout, residual connection and layer norm
        src = self.self_attn_layer_norm(src + self.dropout(_src))

        # src = [batch size, src len, hid dim]

        # positionwise feedforward
        _src = self.positionwise_feedforward(src)

        # dropout, residual and layer norm
        src = self.ff_layer_norm(src + self.dropout(_src))

        # src = [batch size, src len, hid dim]

        return src

Mutli Head Attention Layer

其中一个关键的、新颖的概念是Transformer论文的多头注意力层。
在这里插入图片描述
注意力可以认为是查询向量,键向量和值向量操作的结果，查询向量使用键向量得到注意力向量(通常的输出在softmax操作后，所有的值在0和1之间,总和为1),然后得到一个加权和的值向量。

Transformer 使用缩放点积注意，其中查询和键通过取它们之间的点积结合起来，然后应用softmax操作并缩放 $d_k$ ，最后再乘以值。 $d_k$ 是头维度head_dim，我们将在稍后进一步解释。
在这里插入图片描述
这类似于标准的点积注意，但被缩放了 $d_k$ ，本论文指出，这是用来防止点积的结果变大，导致梯度变得太小。

然而，缩放点积的注意并不是简单地应用于查询、键和值。查询、键和值将它们的hid_dim分割为 $h$ 个头，而不是执行单个注意应用程序，并在所有头上并行计算缩放的点积注意。这意味着，我们将关注 $h$ ，而不是每个注意应用程序关注一个概念。然后，我们将这些头重新组合成它们的hid_dim形状，这样每个hid_dim都有可能关注不同的概念。
在这里插入图片描述
$W^O$ 是应用在多头注意层末尾的线性层fc。 $W^Q, W^K, W^V$ 是线性层fc_q, fc_k和fc_v。

在整个模块中，首先我们用线性层fc_q, fc_k和fc_v计算 $QW^Q$ ， $KW^K$ 和 $VW^V$ ，得到Q, K和V。接下来，我们使用.view将query、key和value的hid_dim分割成n_heads，并正确排列它们，以便它们可以相乘。然后，我们通过将Q和K相乘并将其乘以head_dim的平方根(计算为head_dim = hid_dim // n_heads)来计算energy 能量(未标准化的注意力)。然后我们屏蔽energy 能量，这样我们就不会注意序列中任何我们不应该注意的元素，然后应用softmax和dropout。然后，在将n_heads组合在一起之前，我们将注意力放在value头，V，上。最后，我们乘以这个 $W^O$ ，用fc_o表示。

注意，在我们的实现中，键和值的长度总是相同的，因此，当矩阵乘以softmax的输出，attention，时，和V相乘我们总是有有效的矩阵维数大小。这个乘法是用torch.matmul来实现的,当两个张量都是大于二维时，对每个张量的最后两个维度进行批量矩阵乘法。这将是一个[query_len, key_len] x[key_len，head_dim]的批处理矩阵乘法，它将产生[batch_size，n_heads，query_len，head_dim]的结果。

一开始看起来很奇怪的是，dropout是直接用于注意力。这意味着我们的注意力向量很可能不等于1，我们可能会将全部注意力放在一个标记上，但该标记上的注意力被dropout设置为0。在论文中没有解释，甚至没有提到，但是在官方实现和此后的每个Transformer实现(包括BERT)中都使用了这一点。

class MultiHeadAttentionLayer(nn.Module):
    def __init__(self, hid_dim, n_heads, dropout, device):
        super().__init__()
        
        assert hid_dim % n_heads == 0
        
        self.hid_dim = hid_dim
        self.n_heads = n_heads
        self.head_dim = hid_dim // n_heads
        
        self.fc_q = nn.Linear(hid_dim, hid_dim)
        self.fc_k = nn.Linear(hid_dim, hid_dim)
        self.fc_v = nn.Linear(hid_dim, hid_dim)
        
        self.fc_o = nn.Linear(hid_dim, hid_dim)
        
        self.dropout = nn.Dropout(dropout)
        
        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)
        
    def forward(self, query, key, value, mask = None):
        
        batch_size = query.shape[0]
        
        #query = [batch size, query len, hid dim]
        #key = [batch size, key len, hid dim]
        #value = [batch size, value len, hid dim]
                
        Q = self.fc_q(query)
        K = self.fc_k(key)
        V = self.fc_v(value)
        
        #Q = [batch size, query len, hid dim]
        #K = [batch size, key len, hid dim]
        #V = [batch size, value len, hid dim]
                
        Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)
        
        #Q = [batch size, n heads, query len, head dim]
        #K = [batch size, n heads, key len, head dim]
        #V = [batch size, n heads, value len, head dim]
                
        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale
        
        #energy = [batch size, n heads, query len, key len]
        
        if mask is not None:
            energy = energy.masked_fill(mask == 0, -1e10)
        
        attention = torch.softmax(energy, dim = -1)
                
        #attention = [batch size, n heads, query len, key len]
                
        x = torch.matmul(self.dropout(attention), V)
        
        #x = [batch size, n heads, query len, head dim]
        
        x = x.permute(0, 2, 1, 3).contiguous()
        
        #x = [batch size, query len, n heads, head dim]
        
        x = x.view(batch_size, -1, self.hid_dim)
        
        #x = [batch size, query len, hid dim]
        
        x = self.fc_o(x)
        
        #x = [batch size, query len, hid dim]
        
        return x, attention

Position-wise Feedforward Layer

编码器层内的另一个主要块是position-wise feedforward层，这块相对于多头注意力层来说简单些。输入从hid_dim转换为pf_dim，其中pf_dim通常比hid_dim大很多。原来的Transformer使用hid_dim 512和pf_dim 2048。ReLU激活函数和dropout在被转换回hid_dim表示之前被应用。

为什么使用这个?不幸的是，论文上从来没有解释过。

BERT使用的是GELU激活函数。他们为什么要用 GELU?同样，它也没有得到解释。

class PositionwiseFeedforwardLayer(nn.Module):
    def __init__(self, hid_dim, pf_dim, dropout):
        super().__init__()
        
        self.fc_1 = nn.Linear(hid_dim, pf_dim)
        self.fc_2 = nn.Linear(pf_dim, hid_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        
        #x = [batch size, seq len, hid dim]
        
        x = self.dropout(torch.relu(self.fc_1(x)))
        
        #x = [batch size, seq len, pf dim]
        
        x = self.fc_2(x)
        
        #x = [batch size, seq len, hid dim]
        
        return x

Decoder

解码器的目标是获取源句子 $Z$ 的编码表示，并将其转换为目标句子 $\hat{Y}$ 中的预测标记。然后，我们将 $\hat{Y}$ 与目标语句 $Y$ 中的实际标记进行比较，以计算损失，这将用于计算参数的梯度，然后使用优化器更新权重，以改进预测。
在这里插入图片描述
解码器类似于编码器，但它现在有两个多头注意力层。在目标序列上进行屏蔽的多头注意力层，以及使用解码器表示作为查询，编码器表示作为键和值的多头注意力层。

解码器使用位置嵌入，并将它们与缩放的嵌入目标标记组合，然后进行dropout。同样，我们的位置编码有一个100的“词汇表”，这意味着它们可以接受最长为100个标记的序列。如果需要，可以增加。

组合嵌入然后通过 $N$ 层解码器层，连同编码的源编码（enc_src），源句子掩码和目标句子掩码一起传递。请注意，编码器的层数不必等于解码器的层数，即使它们都用 $N$ 表示。

$N^{th}$ 层之后的解码器表示然后通过一个线性层fc_out。在PyTorch中，softmax操作包含在我们的损失函数中，所以我们不需要在这里显式地使用softmax层。

除了使用源掩码(就像我们在编码器中所做的那样，以防止我们的模型注意到《pad》标记)之外，我们还使用了目标掩码。这将在封装编码器和解码器的Seq2Seq模型中进一步解释，但其要点是，它执行的操作与卷积序列到序列模型中的解码器填充类似。当我们同时处理所有的目标标记时，我们需要一种方法来阻止解码器“作弊”，以只需“查看”目标序列中的下一个标记是什么并输出它。

我们的解码器层也会输出规范化的注意力值，这样我们就可以在稍后绘制它们，看看我们的模型实际上关注的是什么。

class Decoder(nn.Module):
    def __init__(self, output_dim, hid_dim, n_layers, n_heads, pf_dim, dropout, device, max_length=100):
        super().__init__()

        self.device = device

        self.tok_embedding = nn.Embedding(output_dim, hid_dim)
        self.pos_embedding = nn.Embedding(max_length, hid_dim)

        self.layers = nn.ModuleList([DecoderLayer(hid_dim,
                                                  n_heads,
                                                  pf_dim,
                                                  dropout,
                                                  device)
                                     for _ in range(n_layers)])

        self.fc_out = nn.Linear(hid_dim, output_dim)

        self.dropout = nn.Dropout(dropout)

        self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)

    def forward(self, trg, enc_src, trg_mask, src_mask):
        # trg = [batch size, trg len]
        # enc_src = [batch size, src len, hid dim]
        # trg_mask = [batch size, 1, trg len, trg len]
        # src_mask = [batch size, 1, 1, src len]

        batch_size = trg.shape[0]
        trg_len = trg.shape[1]

        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)

        # pos = [batch size, trg len]

        trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))

        # trg = [batch size, trg len, hid dim]

        for layer in self.layers:
            trg, attention = layer(trg, enc_src, trg_mask, src_mask)

        # trg = [batch size, trg len, hid dim]
        # attention = [batch size, n heads, trg len, src len]

        output = self.fc_out(trg)

        # output = [batch size, trg len, output dim]

        return output, attention

Decoder Layer

如前所述，解码器层与编码器层相似，不同之处在于它现在有两个多头注意力层self_attention和encoder_attention。

第一个执行self-attention，就像在编码器中一样，通过使用解码器表示查询、键和值计算和应用注意力。随后是dropout、残差连接和层规范化。这个self_attention层使用目标序列掩码trg_mask，以防止解码器在并行处理目标句子中的所有标记时通过注意它正在处理的标记“前面”的标记来“作弊”。

第二个问题是我们如何将已编码的源语句enc_src实际地输入到解码器中。在这个多头注意力层中，查询是解码器的表示，而键和值是编码器的表示。这里，源掩码src_mask是用来防止多线程注意层注意到源语句中的《pad》标记。然后是dropout层、残差连接层和层规范化层。

最后，我们通过 position-wise feedforward层和另一个dropout的序列，残差连接和层归一化。

解码器层没有引入任何新的概念，只是以略微不同的方式使用与编码器相同的一组层。

class DecoderLayer(nn.Module):
    def __init__(self, hid_dim, n_heads, pf_dim, dropout, device):
        super().__init__()

        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)
        self.enc_attn_layer_norm = nn.LayerNorm(hid_dim)
        self.ff_layer_norm = nn.LayerNorm(hid_dim)
        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
        self.encoder_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)
        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim,
                                                                     pf_dim,
                                                                     dropout)
        self.dropout = nn.Dropout(dropout)

    def forward(self, trg, enc_src, trg_mask, src_mask):
        # trg = [batch size, trg len, hid dim]
        # enc_src = [batch size, src len, hid dim]
        # trg_mask = [batch size, 1, trg len, trg len]
        # src_mask = [batch size, 1, 1, src len]

        # self attention
        _trg, _ = self.self_attention(trg, trg, trg, trg_mask)

        # dropout, residual connection and layer norm
        trg = self.self_attn_layer_norm(trg + self.dropout(_trg))

        # trg = [batch size, trg len, hid dim]

        # encoder attention
        _trg, attention = self.encoder_attention(trg, enc_src, enc_src, src_mask)

        # dropout, residual connection and layer norm
        trg = self.enc_attn_layer_norm(trg + self.dropout(_trg))

        # trg = [batch size, trg len, hid dim]

        # positionwise feedforward
        _trg = self.positionwise_feedforward(trg)

        # dropout, residual and layer norm
        trg = self.ff_layer_norm(trg + self.dropout(_trg))

        # trg = [batch size, trg len, hid dim]
        # attention = [batch size, n heads, trg len, src len]

        return trg, attention

Seq2Seq

最后，我们有一个封装编码器和解码器，以及处理掩码的创建的Seq2Seq模块。

源掩码是通过检查源序列在哪里不等于《pad》标记来创建的。当标记不是《pad》标记时用1标识，当标记是《pad》标记时，用0标识。然后它被unsqueezed成[batch_size，n_heads，seq len, seq len]形状，以便在将掩码应用到energy能量上时能够正确广播。

目标掩码稍微复杂一些。首先，我们为《pad》标记创建一个掩码，就像我们为源掩码所做的那样。接下来，我们使用torch.tril创建一个“subsequent”掩码，trg_sub_mask。这就创建了一个对角线上的元素为零对角线下的元素将被设置为输入张量的值的对角矩阵。在这种情况下，输入张量将是一个满是1的张量。所以这意味着我们的trg_sub_mask看起来像这样(对于一个有5个标记的目标):
在这里插入图片描述
这显示了每个目标标记(行)允许查看的内容(列)。第一个目标标记的掩码为[1,0,0,0]，这意味着它只能查看第一个目标标记。第二个目标标记的掩码为[1,1,0,0,0]，这意味着它可以同时查看第一个和第二个目标标记。

然后，“subsequent”掩码在逻辑上与填充掩码进行处理，这将结合两个掩码，确保后续标记和填充标记都不能被处理。例如，如果最后两个标记是《pad》标记，掩码就会像这样:
在这里插入图片描述
掩码创建后，它们与编码器和解码器以及源句子和目标句子一起使用，以获得我们预测的目标句子，output，以及解码器对源序列的注意。

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder, src_pad_idx, trg_pad_idx, device):
        super().__init__()

        self.encoder = encoder
        self.decoder = decoder
        self.src_pad_idx = src_pad_idx
        self.trg_pad_idx = trg_pad_idx
        self.device = device

    def make_src_mask(self, src):
        # src = [batch size, src len]

        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)

        # src_mask = [batch size, 1, 1, src len]

        return src_mask

    def make_trg_mask(self, trg):
        # trg = [batch size, trg len]

        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)

        # trg_pad_mask = [batch size, 1, 1, trg len]

        trg_len = trg.shape[1]

        trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device=self.device)).bool()

        # trg_sub_mask = [trg len, trg len]

        trg_mask = trg_pad_mask & trg_sub_mask

        # trg_mask = [batch size, 1, trg len, trg len]

        return trg_mask

    def forward(self, src, trg):
        # src = [batch size, src len]
        # trg = [batch size, trg len]

        src_mask = self.make_src_mask(src)
        trg_mask = self.make_trg_mask(trg)

        # src_mask = [batch size, 1, 1, src len]
        # trg_mask = [batch size, 1, trg len, trg len]

        enc_src = self.encoder(src, src_mask)

        # enc_src = [batch size, src len, hid dim]

        output, attention = self.decoder(trg, enc_src, trg_mask, src_mask)

        # output = [batch size, trg len, output dim]
        # attention = [batch size, n heads, trg len, src len]

        return output, attention

训练模型

INPUT_DIM = len(SRC.vocab)
OUTPUT_DIM = len(TRG.vocab)
HID_DIM = 256
ENC_LAYERS = 3
DEC_LAYERS = 3
ENC_HEADS = 8
DEC_HEADS = 8
ENC_PF_DIM = 512
DEC_PF_DIM = 512
ENC_DROPOUT = 0.1
DEC_DROPOUT = 0.1

enc = Encoder(INPUT_DIM, 
              HID_DIM, 
              ENC_LAYERS, 
              ENC_HEADS, 
              ENC_PF_DIM, 
              ENC_DROPOUT, 
              device)

dec = Decoder(OUTPUT_DIM, 
              HID_DIM, 
              DEC_LAYERS, 
              DEC_HEADS, 
              DEC_PF_DIM, 
              DEC_DROPOUT, 
              device)

然后，使用它们来定义我们的整个序列到序列封装模型。

SRC_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]
TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]

model = Seq2Seq(enc, dec, SRC_PAD_IDX, TRG_PAD_IDX, device).to(device)

我们可以检查参数的数量，注意到它明显小于卷积序列对序列模型的37M。

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f'The model has {count_parameters(model):,} trainable parameters')

The model has 9,038,853 trainable parameters

论文没有提及使用的是哪个权值初始化方案，然而Xavier统一似乎在Transformer 模型中很常见，所以我们在这里使用它

def initialize_weights(m):
    if hasattr(m, 'weight') and m.weight.dim() > 1:
        nn.init.xavier_uniform_(m.weight.data)

model.apply(initialize_weights)

在原版Transformer 论文中使用的优化器使用具有“预热”和“冷却”阶段的学习率的Adam。BERT和其他Transformer模型以固定的学习率使用Adam，因此我们将实现这一点。查看这个链接了解更多关于原始Transformer的学习率计划的细节。

注意，学习率需要低于Adam使用的默认值，否则学习是不稳定的。

LEARNING_RATE = 0.0005

optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)

接下来，我们定义我们的损失函数，确保忽略在《pad》标记上计算的损失。

criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)

然后，我们将定义我们的训练循环。这与前一个教程中使用的完全相同。

因为我们想让我们的模型预测《eos》标记，但又不想让它成为我们模型的输入，所以我们只需将《eos》标记从序列的末尾切片。因此:
在这里插入图片描述
$x_i$ 表示实际目标序列元素。然后我们将其输入到模型中，以获得一个预期序列，有望预测标记:

$y_i$ 表示预测的目标序列元素。然后，我们使用原始trg张量计算我们的损失，将《sos》标记前面切掉，留下标记:
在这里插入图片描述
然后我们计算我们的损失并按照标准更新我们的参数。

def train(model, iterator, optimizer, criterion, clip):
    
    model.train()
    
    epoch_loss = 0
    
    for i, batch in enumerate(iterator):
        
        src = batch.src
        trg = batch.trg
        
        optimizer.zero_grad()
        
        output, _ = model(src, trg[:,:-1])
                
        #output = [batch size, trg len - 1, output dim]
        #trg = [batch size, trg len]
            
        output_dim = output.shape[-1]
            
        output = output.contiguous().view(-1, output_dim)
        trg = trg[:,1:].contiguous().view(-1)
                
        #output = [batch size * trg len - 1, output dim]
        #trg = [batch size * trg len - 1]
            
        loss = criterion(output, trg)
        
        loss.backward()
        
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
        
        optimizer.step()
        
        epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

评估循环与训练循环相同，只是没有梯度计算和参数更新。

def evaluate(model, iterator, criterion):
    
    model.eval()
    
    epoch_loss = 0
    
    with torch.no_grad():
    
        for i, batch in enumerate(iterator):

            src = batch.src
            trg = batch.trg

            output, _ = model(src, trg[:,:-1])
            
            #output = [batch size, trg len - 1, output dim]
            #trg = [batch size, trg len]
            
            output_dim = output.shape[-1]
            
            output = output.contiguous().view(-1, output_dim)
            trg = trg[:,1:].contiguous().view(-1)
            
            #output = [batch size * trg len - 1, output dim]
            #trg = [batch size * trg len - 1]
            
            loss = criterion(output, trg)

            epoch_loss += loss.item()
        
    return epoch_loss / len(iterator)

然后我们定义一个小函数，我们可以用它来告诉我们一个epoch循环需要多长时间。

def epoch_time(start_time, end_time):
    elapsed_time = end_time - start_time
    elapsed_mins = int(elapsed_time / 60)
    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
    return elapsed_mins, elapsed_secs

最后，我们训练我们的实际模型。这个模型几乎比卷积序列对序列模型快3倍，也实现了更低的验证复杂性!

N_EPOCHS = 10
CLIP = 1

best_valid_loss = float('inf')

for epoch in range(N_EPOCHS):
    
    start_time = time.time()
    
    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
    valid_loss = evaluate(model, valid_iterator, criterion)
    
    end_time = time.time()
    
    epoch_mins, epoch_secs = epoch_time(start_time, end_time)
    
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'tut6-model.pt')
    
    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
    print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
    print(f'\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')

Epoch: 01 | Time: 0m 11s
	Train Loss: 4.226 | Train PPL:  68.454
	 Val. Loss: 3.046 |  Val. PPL:  21.031
Epoch: 02 | Time: 0m 11s
	Train Loss: 2.824 | Train PPL:  16.837
	 Val. Loss: 2.311 |  Val. PPL:  10.083
Epoch: 03 | Time: 0m 11s
	Train Loss: 2.239 | Train PPL:   9.387
	 Val. Loss: 1.993 |  Val. PPL:   7.341
Epoch: 04 | Time: 0m 11s
	Train Loss: 1.886 | Train PPL:   6.594
	 Val. Loss: 1.814 |  Val. PPL:   6.136
Epoch: 05 | Time: 0m 11s
	Train Loss: 1.640 | Train PPL:   5.156
	 Val. Loss: 1.709 |  Val. PPL:   5.522
Epoch: 06 | Time: 0m 11s
	Train Loss: 1.450 | Train PPL:   4.262
	 Val. Loss: 1.655 |  Val. PPL:   5.234
Epoch: 07 | Time: 0m 11s
	Train Loss: 1.298 | Train PPL:   3.661
	 Val. Loss: 1.631 |  Val. PPL:   5.111
Epoch: 08 | Time: 0m 11s
	Train Loss: 1.171 | Train PPL:   3.226
	 Val. Loss: 1.623 |  Val. PPL:   5.069
Epoch: 09 | Time: 0m 11s
	Train Loss: 1.060 | Train PPL:   2.885
	 Val. Loss: 1.633 |  Val. PPL:   5.119
Epoch: 10 | Time: 0m 11s
	Train Loss: 0.968 | Train PPL:   2.632
	 Val. Loss: 1.637 |  Val. PPL:   5.140

我们加载我们的“最佳”参数，并设法实现比所有以前的模型更好的测试困惑。

model.load_state_dict(torch.load('tut6-model.pt'))

test_loss = evaluate(model, test_iterator, criterion)

print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')

| Test Loss: 1.682 | Test PPL:   5.377 |

推论

现在我们可以使用下面的translate_sentence函数对模型进行翻译。
所采取的步骤是:

如果源语句没有被标记(是字符串)，则标记源语句
附加《sos》和《eos》标记
将源句子数字化
把它转换成张量，然后加上批维数
创建源句子的mask
将源句子和mask输入到编码器中
创建一个列表来保存输出语句，初始化时使用《sos》标记
当我们还没有达到最大长度时
- 将当前输出句子预测转换为具有批维度的张量
- 创建一个目标句子mask
- 将当前输出、编码器输出和两个mask放入解码器
- 从解码器获得下一个输出标记预测以及注意力
- 添加预测到当前输出句子预测
- 如果预测是《eos》标记，则中断循环
将输出语句从索引转换为标记
返回输出语句(删除《sos》标记)和最后一层的注意力

def translate_sentence(sentence, src_field, trg_field, model, device, max_len = 50):
    
    model.eval()
        
    if isinstance(sentence, str):
        nlp = spacy.load('de')
        tokens = [token.text.lower() for token in nlp(sentence)]
    else:
        tokens = [token.lower() for token in sentence]

    tokens = [src_field.init_token] + tokens + [src_field.eos_token]
        
    src_indexes = [src_field.vocab.stoi[token] for token in tokens]

    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)
    
    src_mask = model.make_src_mask(src_tensor)
    
    with torch.no_grad():
        enc_src = model.encoder(src_tensor, src_mask)

    trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]

    for i in range(max_len):

        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)

        trg_mask = model.make_trg_mask(trg_tensor)
        
        with torch.no_grad():
            output, attention = model.decoder(trg_tensor, enc_src, trg_mask, src_mask)
        
        pred_token = output.argmax(2)[:,-1].item()
        
        trg_indexes.append(pred_token)

        if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:
            break
    
    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]
    
    return trg_tokens[1:], attention

现在，我们将定义一个函数，用于显示解码的每个步骤对源句子的注意。因为这个模型有8个头，我们可以看到每个头的注意力。

def display_attention(sentence, translation, attention, n_heads = 8, n_rows = 4, n_cols = 2):
    
    assert n_rows * n_cols == n_heads
    
    fig = plt.figure(figsize=(15,25))
    
    for i in range(n_heads):
        
        ax = fig.add_subplot(n_rows, n_cols, i+1)
        
        _attention = attention.squeeze(0)[i].cpu().detach().numpy()

        cax = ax.matshow(_attention, cmap='bone')

        ax.tick_params(labelsize=12)
        ax.set_xticklabels(['']+['<sos>']+[t.lower() for t in sentence]+['<eos>'], 
                           rotation=45)
        ax.set_yticklabels(['']+translation)

        ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
        ax.yaxis.set_major_locator(ticker.MultipleLocator(1))

    plt.show()
    plt.close()

首先，我们将从训练集获得一个示例。

example_idx = 8

src = vars(train_data.examples[example_idx])['src']
trg = vars(train_data.examples[example_idx])['trg']

print(f'src = {src}')
print(f'trg = {trg}')

src = ['eine', 'frau', 'mit', 'einer', 'großen', 'geldbörse', 'geht', 'an', 'einem', 'tor', 'vorbei', '.']
trg = ['a', 'woman', 'with', 'a', 'large', 'purse', 'is', 'walking', 'by', 'a', 'gate', '.']

我们的翻译看起来很好，尽管我们的模型更改是反复进行的。意思还是一样的。

translation, attention = translate_sentence(src, SRC, TRG, model, device)

print(f'predicted trg = {translation}')

predicted trg = ['a', 'woman', 'with', 'a', 'large', 'purse', 'walks', 'by', 'a', 'gate', '.', '<eos>']

我们可以看到下面每个头的注意力。当然，每一个都是不同的，但很难(也许不可能)去推断大脑到底学会了关注什么。有些头在翻译“a”的时候很注意“eine”，有些头完全不注意，有些头只注意一点。它们似乎都遵循类似的“向下的楼梯”模式，当输出最后两个标记时，注意力平均分布在输入句子中的最后两个标记上。

display_attention(src, translation, attention)

在这里插入图片描述
在验证集上和测试集上的观察注意力的方式同上，此处略过。

BELU

最后计算Transformer的BLEU值。

from torchtext.data.metrics import bleu_score

def calculate_bleu(data, src_field, trg_field, model, device, max_len = 50):
    
    trgs = []
    pred_trgs = []
    
    for datum in data:
        
        src = vars(datum)['src']
        trg = vars(datum)['trg']
        
        pred_trg, _ = translate_sentence(src, src_field, trg_field, model, device, max_len)
        
        #cut off <eos> token
        pred_trg = pred_trg[:-1]
        
        pred_trgs.append(pred_trg)
        trgs.append([trg])
        
    return bleu_score(pred_trgs, trgs)

我们得到的BLEU分数为35.08，高于卷积序列对序列模型的33.3和基于注意力的RNN模型的28.2。所有这一切，同时拥有最少的参数量和最快的训练时间!

bleu_score = calculate_bleu(test_data, SRC, TRG, model, device)

print(f'BLEU score = {bleu_score*100:.2f}')

BLEU score = 35.08

Nlper_Zhao

关注

6
点赞
踩
46

收藏

觉得还不错? 一键收藏
4
评论
Attention is All You Need（Pytorch实现）

在本笔记中，我们将实现一个(稍作修改的版本)的 Attention is All You Need 论文中的transformer模型。本笔记本中的所有图像都将取自transformer 论文。有关Transformer的更多信息，请参见这三篇文章。引言与卷积序列到序列模型类似，Transformer不使用任何递归，它也不使用任何卷积层。相反，该模型完全由线性层、注意力机制和规范化组成。截止到2020年1月，transformer是NLP的主要架构和用于实现许多任务的最先进的结果，并且似乎在不久的将来
复制链接

扫一扫