【模型学习之路】手写+分析Transformer

手写+分析transformer

import torch

a = torch.zeros((666,2,3))
b = torch.zeros((666,3,4))
print((a @ b).shape)   # 666, 2, 4

a = torch.zeros((2,3))
b = torch.zeros((666,3,4))
print((a @ b).shape)  # 666, 2, 4

a = torch.zeros((999,666,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)  # 999, 666, 2, 4

a = torch.zeros((999,1,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)  # 999, 666, 2, 4

a = torch.zeros((666,2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)   # 999, 666, 2, 4

a = torch.zeros((2,3))
b = torch.zeros((999,666,3,4))
print((a @ b).shape)   # 999, 666, 2, 4

多头注意力机制

之前，我们有这样的公式：

就像从一个卷积核向多个卷积核过渡一样，我们可以使用多个注意力头，结合之前提到的多维张量的乘法，得到下面这些式子，其中h是头的数量：

其他的和单头大差不差。

多头注意力层的实现

实现如下图（注意：nn.Linear在应对三维输入时，只会在最后一维运用线性变换，如shape为(p,q,m)的输入经过Linear(m,n)会变成(p,q,n) ，四维及以上同理）。

下图的B显然是Batch_size。

分别指用来生成的矩阵，在这里显然三者都等于输入进来的，这里之所以分开写，是因为后面会出现三者不是同一个矩阵的情况。

其实就是Linear层，这个容易设计。

个人其实感觉多头注意力：

本质上几乎就是“单个更长的注意力头” （事实上，代码里也是这么用的）：

看了一下attention is all you need论文原文，应该是后者，当时写的时候是按前者写的，算了，差别不大。

在最后输出时，transformer还用到了resnet的思想，output加上了输入。

此外，我们在设计网络时，还要考虑一个叫掩码的东东，它的维度和attn（见下图）一样，在特定的场景（模型的具体应用场景）中使用，用来替换或“屏蔽”attn中特定的值，这里也简单实现一下。

先实现Attention函数，四个输入：Q K V mask

def attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
              mask: torch.Tensor = None):
    k = K.size(-1)  # [B, h, m, k]
    scores = Q @ K.transpose(-2, -1) / np.sqrt(k)
    if mask is not None:
        scores = scores.masked_fill(mask, -1e9)
    attn = nn.Softmax(dim=-1)(scores)
    return attn @ V, attn

进一步，实现MultiHeadAttention类：

这里设置了n1和n2，不过在这一步，n1=n2=n，不过在后面会自然会出现两者不一样的情况。

class MultiHeadAttention(nn.Module):
    def __init__(self, n1, n2, h, k, v):
        # n1 表示encoder的， n2 表示decoder的
        super(MultiHeadAttention, self).__init__()
        self.n1 = n1
        self.n2 = n2
        self.h = h
        self.k = k
        self.v = v
        self.W_Q = nn.Linear(n2, h * k, bias=False)
        self.W_K = nn.Linear(n1, h * k, bias=False)
        self.W_V = nn.Linear(n1, h * v, bias=False)
        self.fc = nn.Linear(h * v, n2, bias=False)

    def forward(self, to_Q, to_K, to_V, mask=None):
        res = to_Q
        batch_size = to_Q.size(0)
        Q = self.W_Q(to_Q).view(batch_size, -1, self.h, self.k).transpose(1, 2)
        K = self.W_K(to_K).view(batch_size, -1, self.h, self.k).transpose(1, 2)
        V = self.W_V(to_V).view(batch_size, -1, self.h, self.v).transpose(1, 2)
        if mask is not None:
            mask = mask.unsqueeze(1).repeat(1, self.h, 1)
        out_put, attn = attention(Q, K, V, mask)
        out_put = out_put.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.v)
        out_put = self.fc(out_put) + res
        out_put = nn.LayerNorm(self.n2)(out_put)
        return out_put, attn

我们封装一下，完成Multi-HeadAttention和Add&Norm部分。

Encoder

FeedForwardNet

首先是一个前馈层，这个很简单，一笔带过

class FeedForwardNet(nn.Module):
    def __init__(self, n, d_ff):
        super(FeedForwardNet, self).__init__()
        self.n = n
        self.fc = nn.Sequential(
            nn.Linear(n, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, n, bias=False)
        )

    def forward(self, x):
        out = self.fc(x) + x
        out = nn.LayerNorm(self.n)(out)
        return out

EncoderLayer

然后就可以将我们实现的MultiHeadAttention与前馈层组合起来，形成一个EncoderLayer。

class EncoderLayer(nn.Module):
    def __init__(self, n, h, k, v, d_ff):
        super(EncoderLayer, self).__init__()
        self.d_ff = d_ff
        self.multi_head_attention = MultiHeadAttention(n, n, h, k, v)
        self.feed_forward_net = FeedForwardNet(n, d_ff)

    def forward(self, x, mask=None):
        out, attn = self.multi_head_attention(x, x, x, mask)
        out = self.feed_forward_net(out)
        return out, attn

Encoder

对了，差点忘记写positional encoding的代码了。

以及embedding的代码，这里就不解释embedding是啥意思了。

这里有一次广播，因为不同句子（B个句子（语义块），每个句子m个字，每个字表示为长为n的向量）对应的position矩阵肯定是一样的。

class PositionalEncoding(nn.Module):
    # 这里直接搬运那条知乎作者的代码
    def __init__(self, n, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pos_table = np.array([
            [pos / np.power(10000, 2 * i / n) for i in range(n)]
            if pos != 0 else np.zeros(n) for pos in range(max_len)])
        pos_table[1:, 0::2] = np.sin(pos_table[1:, 0::2])
        pos_table[1:, 1::2] = np.cos(pos_table[1:, 1::2])
        self.pos_table = torch.FloatTensor(pos_table)  # [m, n]

    def forward(self, x):  # x: [B, m, n]
        x += self.pos_table[:x.size(1), :]
        return self.dropout(x)

class Encoder(nn.Module):
    def __init__(self, vocab, n, h, k, v, d_ff, n_layers):
        super(Encoder, self).__init__()
        self.emb = nn.Embedding(vocab, n)  # [B, m] -> [B, m, vocab] -> [B, m, n]
        self.pos = PositionalEncoding(n)
        self.layers = nn.ModuleList([EncoderLayer(n, h, k, v, d_ff)
                                     for _ in range(n_layers)])

    def forward(self, x, mask=None):
        x = self.emb(x)
        x = self.pos(x)
        attn_lst = []
        for layer in self.layers:
            x, attn = layer(x, mask)
            attn_lst.append(attn)
        return x, attn_lst

Decoder

DecoderLayer

首先是单个decoder layer。

这里有两种MultiHeadAttention，一种输入全为decoder里的值，一种还要接收从endoder里出来的值。前者和encoder里面的一个意思，我们重点关注后者。

首先要知道一点，在NLP中，encoder和decoder输入的是不同的语言，也就是说，二者的m,n,vacab都是不一样的。在这里，encoder的表示为m,n。decoder的表示为M,N。这里只是为了写起来方便，代码中用m1 m2 n1 n2。

其中enc_attn里面长这样：

class DecoderLayer(nn.Module):
    def __init__(self, n1, n2, h, k, v, d_ff):
        super(DecoderLayer, self).__init__()
        self.enc_attn = MultiHeadAttention(n1, n2, h, k, v)
        self.dec_attn = MultiHeadAttention(n2, n2, h, k, v)
        self.ffn = FeedForwardNet(n2, d_ff)

    def forward(self, dec_in, enc_out, enc_mask=None, dec_mask=None):
        out, dec_attn = self.dec_attn(dec_in, dec_in, dec_in, dec_mask)
        out, enc_attn = self.enc_attn(out, enc_out, enc_out, enc_mask)
        out = self.ffn(out)
        return out, dec_attn, enc_attn

Decoder

然后做整个Decoder

class Decoder(nn.Module):
    def __init__(self, vocab, n, h, k, v, d_ff, n_layers):
        super(Decoder, self).__init__()
        self.emb = nn.Embedding(vocab, n)
        self.pos = PositionalEncoding(n)
        self.layers = nn.ModuleList([DecoderLayer(n, h, k, v, d_ff)
                                     for _ in range(n_layers)])

    def forward(self, x, enc_out, enc_mask=None, dec_mask=None):
        x = self.emb(x)
        x = self.pos(x)
        attn_lst = []
        for layer in self.layers:
            x, dec_attn, enc_attn = layer(x, enc_out, enc_mask, dec_mask)
            attn_lst.append((dec_attn, enc_attn))
        return x, attn_lst

组装Trasformer!

最后，组装！

class Transformer(nn.Module):
    def __init__(self, vocab1, vocab2, n1, n2, h, k, v, d_ff, n_layers,
                 enc_mask=None, dec_mask=None, enc_dec_mask=None):
        super(Transformer, self).__init__()
        self.enc_mask = enc_mask
        self.dec_mask = dec_mask
        self.enc_dec_mask = enc_dec_mask
        self.encoder = Encoder(vocab1, n1, h, k, v, d_ff, n_layers)
        self.decoder = Decoder(vocab2, n1, n2, h, k, v, d_ff, n_layers)
        self.out = nn.Linear(n2, vocab2, bias=False)  # [B, m, n] -> [B, m, vocab]
        
    def forward(self, enc_in, dec_in):
        enc_out, enc_attn_lst = self.encoder(enc_in, self.enc_mask)
        dec_out, dec_attn_lst = self.decoder(dec_in, enc_out, self.enc_dec_mask, self.dec_mask)
        out = self.out(dec_out)
        return out, enc_attn_lst, dec_attn_lst

后话

测试一下

if __name__ == '__main__':
    test_enc_in = torch.randint(0, 50, (2, 10))  # [B, m1]
    test_dec_in = torch.randint(0, 50, (2, 12))  # [B, m2]

    model = Transformer(vocab1=100, vocab2=70, n1=32, n2=64, h=8, k=32, v=128, d_ff=128, n_layers=6,
                        enc_mask=None, dec_mask=None, enc_dec_mask=None)
    test_out, enc_attn_lst, dec_attn_lst = model(test_enc_in, test_dec_in)
    print(test_out.shape)  # [B, m2, vocab2]
    # torch.Size([2, 12, 70])