transformer pytorch transformer pytorch 量化

本文主要根据“Attention Is All You Need”里的提到的transformer来实现的。
主要参考了：
http://nlp.seas.harvard.edu/2018/04/03/attention.htmlhttps://kexue.fm/archives/4765

概述

在过去的一年中，根据“Attention Is Al You Need”所提到的transformer已经给很多人留下了深刻的印象。除了在翻译质量方面取得重大进步外，它还为许多其他NLP任务提供了新的架构。论文本身写得非常清楚，但传统观点认为很难正确实施。
在这篇文章中，我以逐行实现的形式呈现了论文的“注释”版本。我重新排序并删除了原始论文中的一些部分，并在整个过程中添加了评论。这个文档本身就是一个有效的笔记本，应该是一个完全可用的实现。总共有400行代码，可以在4个GPU上每秒处理27,000个指令。
对于模型的其他架构实现[Tensor2Tensor]（tensorflow）和Sockeye（mxnet）。

背景

依赖库：

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
from torch.autograd import Variable
import matplotlib.pyplot as plt
import seaborn
seaborn.set_context(context="talk")

为减少序列计算，形成了扩展GPU，ByteNet和ConvS2S等基础神经网络，这些网络使用卷积神经网络作为基本构建块，并行计算所有输入和输出位置的隐藏层。在这些模型中，关联来自两个任意输入或输出位置的信号所需的操作数量随着位置距离的增加而增加，对于ConvS2S呈线性增长，对于ByteNet呈对数增长。这使得学习远程位置之间的依赖性变得更加困难。而在transformer中，这被减少到恒定的操作次数，尽管由于Attention加权平均而导致有效分辨率降低，但我们Multi-Head Attention来抵消这种损失。
self-Attention，有时称为intra-attention是一种Attention机制，其关联单个序列的不同位置以计算序列的特征。self-Attention已经成功地用于各种任务，包括阅读理解，抽象概括，文本蕴涵和学习任务独立的句子表示。端到端记忆网络基于循环注意Attention而不是循环对齐序列，并且已经证明在简单语言问答和语言建模任务上表现良好。
而据我们所知，transformer是第一个完全依靠self-Attention的转换模型来计算其输入和输出的特征，而不使用序列对齐的RNN或卷积。

模型结构

大多数seq模型具有encoder - decoder结构, 这里，encoder将符号表示的输入序列$（x_1，...，x_n）$映射到连续表示序列$\mathbf {z} =（z_1，...，z_n）$。给定$\mathbf {z}$，decoder然后一次生成一个元素的输出序列$（y_1，...，y_m）$。在每个步骤中，模型是自动回归$（x_1，...，x_n）$到一系列连续表示$\mathbf {z} =（z_1，...，z_n）$。给定$\mathbf {z}$，decoder然后一次一个元素地生成符号的输出序列$（y_1，...，y_m）$。在每个步骤中，模型都是自动回归的，在生成下一个输出，将先前生成的符号作为附加输入。
Encoder-Decoder:

class EncoderDecoder(nn.Module):
    """
    A standard Encoder-Decoder architecture. Base for this and many 
    other models.
    """
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
        
    def forward(self, src, tgt, src_mask, tgt_mask):
        "Take in and process masked src and target sequences."
        return self.decode(self.encode(src, src_mask), src_mask,
                            tgt, tgt_mask)
    
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
    
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

class Generator(nn.Module):
    "Define standard linear + softmax generation step."
    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        return F.log_softmax(self.proj(x), dim=-1)

transformer遵循这种整体架构，使用堆叠的self-Attention（Multi-Head Attention），并全连接层（Feed Forward）用于编码器和解码器，分别如图1的左半部分和右半部分所示。

Encoder and Decoder Stacks

Encoder

Encoder由N = 6个相同层的Stacks组成。
克隆模型块，克隆的模型块参数不共享

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

我们在两个子层中的输出后使用残余连接，然后是层标准化。也就是说，每个子层的输出是$LayerNorm(x + Sublayer(x))$，其中$Sublayer(x)$是由子层本身实现的功能。我们将dropout应用于每个子层的输出，然后将其添加到子层输入并进行规范化。为了方便参差连接，模型中的所有子层以及embedding 层都产生维度为512维的输出。

class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(sublayer(self.norm(x)))

每层有两个子层。 第一种是Multi-Head Attention机制，第二种是简单的，逐点全连接的feed- forward 前馈网络。

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

Decoder

Decoder也由N = 6个相同层的Stacks组成。

class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

除了每个encoder层中的两个子层之外，decoder还插入第三子层，其对encoder的输出执行Multi-Head Attention。相似于encoder, 我们在每个子层后使用残余连接，然后进行层标准化。

class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)
 
    def forward(self, x, memory, src_mask, tgt_mask):
        "Follow Figure 1 (right) for connections."
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

我们还修改decoder堆栈中的self-Attention子层以防止位置出现在后续位置。确保了位置i的预测仅依赖于小于i的位置处的已知输出。不知道这么翻译看原文：
We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking, combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i.

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    return torch.from_numpy(subsequent_mask) == 0

就是在decoder层的self-Attention中，由于生成$s_i$时，$s_{i+1}$并没有产生，所以不能有$s_i$和$s_{i+1}$的关联系数，即只有下三角矩阵有系数，即上图黄色部分。

Attention

Attention可以被描述为将query和key-value映射到输出，其中query, keys, values和输出都是向量。输出被计算为value的加权和，其中分配给每个值的权重由query与key的兼容性函数计算。详细了解Attention可以参考这篇博客 其中一种Attention叫“Scaled Dot-Product Attention”。输入包含维度$d_k$的query和key，维度$d_v$的value。我们计算query与所有key的矩阵乘积，并除以$\sqrt{d_K}$，然后使用$softmax(x)$函数获取对应value的权重$a_{ij}$。整个过程如下图所示：

在实际操作中，我们同时在一组query上计算Attention函数，将它们打包成矩阵$Q$。key向量组合成矩阵$K$和value向量组合成矩阵$V$，从而可以计算输出通过以下矩阵运算：

$Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V$

实现代码如下：

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

两个最常用的Attention是additive attention和dot-product (multiplicative) attention。dot-product attention除了除以$\frac{1}{\sqrt{d_k}}$，其他一致。Additive attention使用具有单个隐藏层的feed-forward网络来计算query与key的关联程度。虽然两者在理论上的复杂性相似，但在实践中，dot-product attention更快，更节省空间，因为它可以使用高度优化的矩阵乘法来实现。
而对于较小的$d_k$值，两种机制的表现相似，但additive attention比dot-product attention表现更好，没有缩放为更大的$d_k$值。我们怀疑对于更大的$d_k$值，dot-product attention增加更大，使得softmax函数进入具有极小梯度的区域。（为了说明点积变大的原因，假设$q$和$k$的分量是独立的随机变量，均值为0，方差为1。然后他们进行点积，$q \bullet k= \sum_{i=1}^{n}q_ik_i$ 则有均值0和方差$d_k$）。因此为了抵消这种影响，我们在点乘的基础上除以$\frac{1}{\sqrt{d_k}}$。
而对于Multi-head attention允许模型共同关注来自不同位置的不同表示子空间的信息。要避免上诉情况不是除以$\frac{1}{\sqrt{d_k}}$而是要乘以矩阵做标准化，$MulitHead(Q,K,V)=Concat（head_1,......head_h）W^0$

$head_i = Attention(QW_{i}^{k},KW_{i}^{k},VW_{i}^{k},)$

其中$QW_{i}^{k}$，表示向量$Q$的一种投影，其中$W$为投影参数，$W_{i}^{Q}\in R^{d_{model}d_k}$,$W_{i}^{K}\in R^{d_{model}d_k}$,$W_{i}^{V}\in R^{d_{model}d_v}$和$W^{0}\in R^{d_{model}\times hd_v}$。在本文中，我们采用h = 8个平行attention层。对于每一个Attention我们使用$dk=dv=dmodel/h=64$。由于每个头部的尺寸减小，总计算成本与具有全维度的单头Attention相似。计算流程图如下：

代码如下：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        
        # 1) Do all the linear projections in batch from d_model => h x d_k 
        query, key, value = \
            [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
             for l, x in zip(self.linears, (query, key, value))]
        
        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask, 
                                 dropout=self.dropout)
        
        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
             .view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

Attention在我们的模型中的应用

Transformer以三种不同的方式使用Multi-head attention：
1）在“encoder-decoder attention”层中，query来自先前的解码器层，并且存储器key和value来自编码器的输出。这允许解码器中的每个位置都参与输入序列中的所有位置。这模仿了seq-seq模型中的典型encoder-decoder attention机制。
2） encoder层包含 self-attention 层，在 self-attention 层，所有的query, key, value来自同一个位置，在这种情况下，是编码器中前一层的输出。编码器中的每个位置都可以处理编码器前一层中的所有位置。
3）类似地，解码器中的自注意层允许解码器中的每个位置参与解码器中的所有位置直到并包括该位置。我们需要防止解码器中的向左信息流以保持自回归属性。我们通过屏蔽softmax输入中与非法连接相对应的所有值来实现缩放的dot-product Attention。

Position-wise Feed-Forward Networks

除了Attention子层，我们的编码器和解码器中的每个层都包含一个完连接的前馈网络，该网络分别和相同地应用于每个位置。这包括两个线性变换，其间有ReLU激活。
$FFN(x)=max(0,xW1+b1)W2+b2$
虽然线性变换在不同位置上是相同的，但它们在层与层之间使用不同的参数。另一种描述这种情况的方法是两个内核大小为1的卷积。输入和输出的维数是$d_{model} = 512$，内层的维度$dff = 2048$。

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

Embeddings and Softmax

与其他序列传导模型类似，我们使用learned embeddings将输入标记和输出标记转换为维度$d_{model}$的向量。我们还使用通常学习的线性变换和softmax函数将解码器输出转换为预测的下一个标签的概率。在我们的模型中，我们在两个embedding层和pre-softmax线性变换层之间共享相同的权重矩阵。在embedding层中，我们将这些权重乘以$\sqrt{d_{model}}$。

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

位置编码

截止目前为止，我们介绍的Transformer模型并没有捕捉顺序序列的能力，也就是说无论句子的结构怎么打乱，Transformer都会得到类似的结果。换句话说，Transformer只是一个功能更强大的词袋模型而已。
为了解决这个问题，论文中在编码词向量时引入了位置编码（Position Embedding）的特征。具体地说，位置编码会在词向量中加入了单词的位置信息，这样Transformer就能区分不同位置的单词了。
那么怎么编码这个位置信息呢？常见的模式有：a. 根据数据学习；b. 自己设计编码规则。在这里作者采用了第二种方式。那么这个位置编码该是什么样子呢？通常位置编码是一个长度为 $d_{model}$ 的特征向量，这样便于和词向量进行单位加的操作。编码公式如下：
$PE_{(pos,2i)}=sin(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$
$PE_{(pos,2i+1)}=cos(\frac{pos}{10000^{\frac{2i}{d_{model}}}})$
其中$pos$ 表示单词的位置，$i$表示单词的维度。也就是说，位置编码的每个维度对应于正弦曲线。波长形成从$2π$到$10000⋅2π$的几何级数。我们之所以选择这个函数，是因为我们假设它可以让模型轻松地通过相对位置来学习，因为对于任何固定偏移量$k$，$PE_{pos + k}$可以表示为$PE_{pos}$的线性函数。此外，我们将dropout应用于编码器和解码器堆栈中的embedding和位置编码中。对于基本模型，我们使用$P_{drop} = 0.1$的速率。

class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 用了个技巧先计算log的在计算exp
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        # position * div_term 这里生成一个以pos为行坐标，i为列坐标的矩阵
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        x = x + Variable(self.pe[:, :x.size(1)], 
                         requires_grad=False)  # x.size(1)就是有多少个pos
        return self.dropout(x)

代码不难，多看几遍就理解了

整个模型

在这里，我们定义一个函数，它接受超参数并生成一个完整的模型。

def make_model(src_vocab, tgt_vocab, N=6, 
               d_model=512, d_ff=2048, h=8, dropout=0.1):
    "Helper: Construct a model from hyperparameters."
    c = copy.deepcopy
    attn = MultiHeadedAttention(h, d_model)
    ff = PositionwiseFeedForward(d_model, d_ff, dropout)
    position = PositionalEncoding(d_model, dropout)
    model = EncoderDecoder(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        Decoder(DecoderLayer(d_model, c(attn), c(attn), 
                             c(ff), dropout), N),
        nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
        nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
        Generator(d_model, tgt_vocab))
    
    # This was important from their code. 
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model

Train

本节介绍了我们模型的训练方式。
我们快速介绍一下训练标准编码器解码器模型所需的一些工具。 首先，我们定义一个批处理对象，它包含用于训练的src和目标句子，以及构造掩码。

class Batch:
    "Object for holding a batch of data with mask during training."
    def __init__(self, src, trg=None, pad=0):
        self.src = src
        self.src_mask = (src != pad).unsqueeze(-2)
        if trg is not None:
            self.trg = trg[:, :-1]
            self.trg_y = trg[:, 1:]
            self.trg_mask = \
                self.make_std_mask(self.trg, pad)
            self.ntokens = (self.trg_y != pad).data.sum()
    
    @staticmethod
    def make_std_mask(tgt, pad):
        "Create a mask to hide padding and future words."
        tgt_mask = (tgt != pad).unsqueeze(-2)
        tgt_mask = tgt_mask & Variable(
            subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))
        return tgt_mask

接下来，我们创建一个通用的训练和评分功能来跟踪损失。 我们传入一个通用的损失计算函数，它也处理参数更新。

def run_epoch(data_iter, model, loss_compute):
    "Standard Training and Logging Function"
    start = time.time()
    total_tokens = 0
    total_loss = 0
    tokens = 0
    for i, batch in enumerate(data_iter):
        out = model.forward(batch.src, batch.trg, 
                            batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)
        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens
        if i % 50 == 1:
            elapsed = time.time() - start
            print("Epoch Step: %d Loss: %f Tokens per Sec: %f" %
                    (i, loss / batch.ntokens, tokens / elapsed))
            start = time.time()
            tokens = 0
    return total_loss / total_tokens