Transformer代码实战_transformer 源码 博客 github

1. 本文适合那些对于Transformer了解，但不清楚里面的矩阵运算的维度，以及一些细节实现，比如decoder的自掩码，输入字符的Embedding，Q的线性映射的细节，代码中都有详细的注释。
2. 本文的代码借鉴自：源码
   以及一个视频讲解：Transformer代码(源码Pytorch版本)从零解读(Pytorch版本）
   我对一些地方做了修改，并添加了中文注释
3. 建议先看一遍源代码，调试一遍，不懂的看注释，最后还有些困惑的再看看我的手写稿，里面以一个具体的例子描述了Transformer的全过程。

# 本代码借鉴自以下地址，加了中文注释，并修改了我认为不合理的地方
# https://github.com/graykode/nlp-tutorial/blob/master/5-1.Transformer/Transformer.ipynb
import numpy as np
import torch
import torch.nn as nn


# 将输入的句子用自己定义的简单词典编码
def make_batch(sentences):
    # 注意这里是两层括号，也就是比如input_batch的大小是：1*4
    # Transformer中训练是以句子为单位，就和图像分割中训练是以图像为单位一样，我们这里仅用1个句子来模拟，batch_size = 1
    input_batch = [[src_vocab[n] for n in sentences[0].split()]]
    decoder_input_batch = [[tgt_vocab[n] for n in sentences[1].split()]]
    target_batch = [[tgt_vocab[n] for n in sentences[2].split()]]
    return torch.LongTensor(input_batch), torch.LongTensor(decoder_input_batch), torch.LongTensor(target_batch)


def get_attn_pad_mask(seq_q, seq_k):
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    # 词表中Pad的设置为0
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # batch_size x 1 x len_k(=len_q), one is masking
    # 返回一个长度为batch_size x len_q x len_k的矩阵，其中pad的位置用True表示，避免pad影响计算注意力得分
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # batch_size x len_q x len_k


"""
    这里的n_position是输入句子的最大长度（比如我们的例子中输入是“我爱你P”，n_position=4）
    每一个词首先都会被one-hot编码（这里为了省事用的是顺序编码但这是不对的），顺序编码后“我爱你P”的编码对照词表为“1230”，输入的数据就表示为[[1, 2, 3, 0]]，这里两个括号是batch_size=1
    接着我们通过一个nn.embedding层，将顺序编码（1位表示）转变为词嵌入（d_model=512位表示），也就是说输入从1*4转换为了1*4*512，一个词用512位表示
    因此，对于第一个词的位置0，还需要遍历d_model=512，求这个词的实际编码的各个位置，也就是get_immutable_vector（）函数
    最后返回的位置编码也是一个1*4*512的矩阵，表达了每个词编码的位置
"""
def get_sinusoid_encoding_table(n_position, d_model):
    # 这个内部函数获取的position是输入句子中一个词的position，
    def get_immutable_vector(position):
        immutable_vector = []
        for i in range(d_model):
            # 这里原代码是i // 2，论文是 i * 2，我选择了论文的做法
            immutable_vector.append(position / np.power(10000, 2 * i / d_model))
        return immutable_vector

    sinusoid_table = np.array([get_immutable_vector(every_position) for every_position in range(n_position)])
    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # 上面分析了位置编码的矩阵为4*512，比如对于第一个词的编码，偶数位使用sin变换
    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])
    return torch.FloatTensor(sinusoid_table)

def get_attn_subsequent_mask(seq):
    """
    这个方法也写的很巧妙，对于target: 'S I love you', 第一次看见S，看不到后面，第二次看到S I，看不到后面
    其实就是一个上三角都是1的矩阵
    0 1 1 1
    0 0 1 1
    0 0 0 1
    0 0 0 0
    每次能看到的就是每一行的0代表的序列
    """
    # seq：[batch_size, tgt_len, d_model]
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    subsequent_mask = np.triu(np.ones(attn_shape), k=1)
    subsequent_mask = torch.from_numpy(subsequent_mask).byte()
    return subsequent_mask

class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        scores = torch.matmul(Q, K.transpose(-1, -2))
        scores.masked_fill_(attn_mask, 1e-9)
        attn = nn.Softmax(dim=-1)(scores)
        context = torch.matmul(attn, V)
        return context


class FeedForwardNet(nn.Module):
    def __init__(self):
        super(FeedForwardNet, self).__init__()
        # 论文中有说，两个线性层和一个conv1d效果一样
        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, enc_inputs):
        # enc_inputs: [batch * len_q * d_model]
        residual = enc_inputs
        enc_inputs = nn.ReLU()(self.conv1(enc_inputs.transpose(1, 2)))
        enc_inputs = self.conv2(enc_inputs).transpose(1, 2)
        return self.layer_norm(enc_inputs + residual)

class MultiHeadAttn(nn.Module):
    def __init__(self):
        super(MultiHeadAttn, self).__init__()
        self.W_Q = nn.Linear(d_model, d_k * n_heads)
        self.W_K = nn.Linear(d_model, d_k * n_heads)
        self.W_V = nn.Linear(d_model, d_v * n_heads)
        self.linear = nn.Linear(n_heads * d_v, d_model)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, Q, K, V, attn_mask):
        residual, batch_size = Q, Q.size(0)
        # Q: [batch_size x len_q x d_model], K: [batch_size x len_k x d_model], V: [batch_size x len_k x d_model]
        # 对于encoder层，自注意力：输入Q=K=V，大小1*4*512
        # W_Q(Q)，就是一个线性映射，从512个输入映射到d_k*n_heads=512（这两个都是自定义参数）个输出
        # 注意这里  不同于  线性回归，对于输入Q=K=V（1*4*512），线性回归，是把矩阵的每一行，当做一个样本，然后y=Xw (4*513  513*1 = 4*1)，得到每个样本经过神经元后的输出，然后求损失函数这样
        # 这里的Q（1*4*512）我们，进行W_Q(Q)只是把他丢进去做一个线性变换输出的结果是（1*4*（d_k*n_heads））
        q_s = self.W_Q(Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # q_s: [batch*n_heads*len_q*d_k]
        # 至于后面为什么要这样切(view函数)再转置，而不是直接像下面这样切
        # q_s = self.W_Q(Q).view(batch_size, n_heads, -1, d_k)  # q_s: [batch*n_heads*len_q*d_k]
        # print(torch.allclose(q_s, q_t, rtol=1e-5))  # False
        #   view()中，第一刀先把整个数据分成各个batch_size肯定没问题
        #   然后第二刀，其实是在分割各个单词，也就是说多头的q_i, k_i, v_i来自的是一个句子的每个单词的某些部分，比如“我爱你P”，每个单词由512个数字表示，分8头，第一头就是“我”里取0-64位+“爱”里取0-64位+“你”里取0-64位+“P”里取0-64位
        #   最后的d_k这一刀，则是为了保证QK之间相乘时不报错
        #   而第二个形式的分割就是：“我爱你P”，第一个头“我”里面取512个，也就是第一个头就只处理“我”这个单词，并不是我们想要的，论文中想要的是多头自注意力能够从多个角度（子空间）看待问题

        k_s = self.W_K(K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # k_s: [batch*n_heads*len_k*d_k]
        v_s = self.W_V(V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)  # v_s: [batch*n_heads*len_k*d_v]

        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)  # attn_mask : [batch_size x n_heads x len_q x len_k]

        # context: [batch_size x n_heads x len_q x d_v], attn: [batch_size x n_heads x len_q(=len_k) x len_k(=len_q)]
        context = ScaledDotProductAttention()(q_s, k_s, v_s, attn_mask)
        context = context.transpose(1, 2).contiguous().view(batch_size, -1, n_heads * d_v)  # context: [batch_size x len_q x n_heads x d_v]
        output = self.linear(context)
        return self.layer_norm(output + residual)  # output: [batch_size x len_q x d_model]


class Encoderlayer(nn.Module):
    def __init__(self):
        super(Encoderlayer, self).__init__()
        self.enc_multi_head_attn = MultiHeadAttn()
        self.ffn = FeedForwardNet()

    def forward(self, enc_inputs, enc_self_attn_mask):
        context= self.enc_multi_head_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)
        enc_outputs = self.ffn(context)
        return enc_outputs

class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttn()
        self.dec_enc_attn = MultiHeadAttn()
        self.pos_ffn = FeedForwardNet()

    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        dec_outputs = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)
        dec_outputs = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)
        dec_outputs = self.pos_ffn(dec_outputs)
        return dec_outputs



class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)
        # 这一行首先nn.Embedding.from_pretrained()就是用一个给定的词向量索引矩阵，来转换输入的文本
        # src_vocab_size代表为每一个输入词典中的词，每个词用d_model个数字表示，每个数字，都生成一个位置索引
        # 这一行比较复杂，下面有等价的写法（forward里），比较好理解
        self.pos_emb = nn.Embedding.from_pretrained(get_sinusoid_encoding_table(src_vocab_size, d_model), freeze=True)
        self.layers = nn.ModuleList([Encoderlayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):
        # 这里写的真的非常巧妙，其实就是定义了一个每个编码位置的索引矩阵：1*4*512（batch_size*src_vocab*d_model），根据论文中的方法给每个词的每一个位置一个编码
        enc_outputs = self.src_emb(enc_inputs) + self.pos_emb(enc_inputs)
        # 也可以这样写，更容易理解：
        """
        # 效果完全等价，读者可以自行调试尝试，但原作者的写法真的非常高级，很巧妙的利用了nn.embedding函数，可见功力之深
        position_vector = get_sinusoid_encoding_table(src_vocab_size, d_model)  # torch.size()为4*512
        # 获取每个词的位置信息：
        pos_emb = []
        for idx in enc_inputs:   # [1, 2, 3, 0]
            pos_emb.append(position_vector[idx])
        pos_emb = torch.cat(pos_emb, dim=0).view(-1, 512)
        enc_outputs = self.src_emb(enc_inputs) + pos_emb
        """
        enc_self_attn_pad_mask = get_attn_pad_mask(enc_inputs, enc_inputs)
        for layer in self.layers:
            enc_outputs = layer(enc_outputs, enc_self_attn_pad_mask)  # enc_outputs: [batch_size x len_q x d_model]
        return enc_outputs


class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = nn.Embedding.from_pretrained(get_sinusoid_encoding_table(tgt_vocab_size, d_model), freeze=True)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])
    def forward(self, dec_inputs, enc_outputs, enc_inputs):
        dec_outputs = self.tgt_emb(dec_inputs) + self.pos_emb(dec_inputs)  # dec_inputs: [batch_size * tgt_len * d_model]
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs)  # dec_self_attn_pad_mask: [batch_size * tgt_len * tgt_len]
        dec_self_attn_subsequent_mask = get_attn_subsequent_mask(dec_inputs)
        dec_self_attn_mask = torch.gt(dec_self_attn_pad_mask+dec_self_attn_subsequent_mask, 0)

        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)

        for layer in self.layers:
            dec_outputs = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)
        return dec_outputs


class Transformer(nn.Module):
    """
    编写init函数，先从大局看
    整体看Transformer分为编码，解码，以及结果输出三个模块
    于是分别定义如下
    """
    def __init__(self):
        super(Transformer, self).__init__()
        self.encoder = Encoder()
        self.decoder = Decoder()
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False)

    def forward(self, enc_inputs, dec_input):
        enc_outputs = self.encoder(enc_inputs)
        dec_outputs = self.decoder(dec_input, enc_outputs, enc_inputs)
        dec_logits = self.projection(dec_outputs)
        return dec_logits.view(-1, dec_logits.size(-1))

if __name__ == '__main__':
    sentences = ['我 爱 你 P', 'S I love you', 'I love you E']

    # 原句词典
    src_vocab = {'P': 0, '我': 1, '爱': 2, '你': 3}
    src_vocab_size = len(src_vocab)

    # 翻译后的句子词典
    tgt_vocab = {'P': 0, 'I': 1, 'love': 2, 'you': 3, 'S': 4, 'E': 5}
    tgt_vocab_size = len(tgt_vocab)
    number_dict = {i: w for i, w in enumerate(tgt_vocab)}

    src_len = 4
    tgt_len = 4

    d_model = 512  # 输入编码的词向量的长度，以及模型在运行过程中的各个层输入输出向量的维度
    d_ff = 2048  # 2层全连接层的中间维度
    d_k = d_v = 64  # 设置的q=k, v矩阵线性转换后的维度
    n_layers = 6  # 编码、解码层重复的次数
    n_heads = 8  # 几头注意力机制
    model = Transformer()

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    enc_inputs, dec_inputs, target = make_batch(sentences)

    for epoch in range(20):
        optimizer.zero_grad()
        outputs = model(enc_inputs, dec_inputs)
        loss = criterion(outputs, target.contiguous().view(-1))
        print('[TRAIN] epoch：', '%04d' % (epoch + 1), ', cost：', '%.6f' % loss)
        loss.backward()
        optimizer.step()

    # Test
    predict = model(enc_inputs, dec_inputs)
    predict = predict.data.max(1, keepdim=True)[1]
    print(sentences[0], '->', [number_dict[n.item()] for n in predict.squeeze()])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259

• 本文的github地址：https://github.com/rainyNighti/Implement_CV_code_using_Pytorch/blob/master/Transformer/Transformer.py
• 手写推导过程
  • 一些参数：

batch_size=2
max_src_sentence_len=6
max_tgt_sentence_len=8
# 模拟输入两句话：“今天天气好P”，“我爱你PPP”
# 填充符号：P:0
# 翻译的结果：“Nice weather today E P P P P”，“I love you E P P P P”
# decoder_input：“S Nice weather today P P P P”，“S I love you P P P P”
1
2
3
4
5
6
7