Transformer的PyTorch实现
Google 2017年的論文 Attention is all you need 闡釋了什么叫做大道至簡！該論文提出了Transformer模型，完全基于Attention mechanism，拋棄了傳統(tǒng)的RNN和CNN。
我們根據(jù)論文的結(jié)構(gòu)圖，一步一步使用 PyTorch 實現(xiàn)這個Transformer模型。
Transformer架構(gòu)
首先看一下transformer的結(jié)構(gòu)圖：
解釋一下這個結(jié)構(gòu)圖。首先，Transformer模型也是使用經(jīng)典的encoer-decoder架構(gòu)，由encoder和decoder兩部分組成。
上圖的左半邊用Nx框出來的，就是我們的encoder的一層。encoder一共有6層這樣的結(jié)構(gòu)。
上圖的右半邊用Nx框出來的，就是我們的decoder的一層。decoder一共有6層這樣的結(jié)構(gòu)。
輸入序列經(jīng)過word embedding和positional encoding相加后，輸入到encoder。
輸出序列經(jīng)過word embedding和positional encoding相加后，輸入到decoder。
最后，decoder輸出的結(jié)果，經(jīng)過一個線性層，然后計算softmax。
word embedding和positional encoding我后面會解釋。我們首先詳細地分析一下encoder和decoder的每一層是怎么樣的。
Encoder
encoder由6層相同的層組成，每一層分別由兩部分組成：
第一部分是一個multi-head self-attention mechanism
第二部分是一個position-wise feed-forward network，是一個全連接層
兩個部分，都有一個 殘差連接(residual connection)，然后接著一個Layer Normalization。
如果你是一個新手，你可能會問：
multi-head self-attention 是什么呢？
參差結(jié)構(gòu)是什么呢？
Layer Normalization又是什么？
這些問題我們在后面會一一解答。
Decoder
和encoder類似，decoder由6個相同的層組成，每一個層包括以下3個部分：
第一個部分是multi-head self-attention mechanism
第二部分是multi-head context-attention mechanism
第三部分是一個position-wise feed-forward network
還是和encoder類似，上面三個部分的每一個部分，都有一個殘差連接，后接一個Layer Normalization。
但是，decoder出現(xiàn)了一個新的東西multi-head context-attention mechanism。這個東西其實也不復雜，理解了multi-head self-attention你就可以理解multi-head context-attention。這個我們后面會講解。
Attention機制
在講清楚各種attention之前，我們得先把attention機制說清楚。
通俗來說，attention是指，對于某個時刻的輸出y，它在輸入x上各個部分的注意力。這個注意力實際上可以理解為權(quán)重。
attention機制也可以分成很多種。Attention? Attention! 一文有一張比較全面的表格：
Figure 2. a summary table of several popular attention mechanisms.
上面第一種additive attention你可能聽過。以前我們的seq2seq模型里面，使用attention機制，這種**加性注意力(additive attention)**用的很多。Google的項目 tensorflow/nmt 里面使用的attention就是這種。
為什么這種attention叫做additive attention呢？很簡單，對于輸入序列隱狀態(tài)hi h_ih
i
?
和輸出序列的隱狀態(tài)st s_ts
t
?
，它的處理方式很簡單，直接合并，變成[st;hi] [s_t;h_i][s
t
?
;h
i
?
]
但是我們的transformer模型使用的不是這種attention機制，使用的是另一種，叫做乘性注意力(multiplicative attention)。
那么這種乘性注意力機制是怎么樣的呢？從上表中的公式也可以看出來：兩個隱狀態(tài)進行點積！
Self-attention是什么？
到這里就可以解釋什么是self-attention了。
上面我們說attention機制的時候，都會說到兩個隱狀態(tài)，分別是hi h_ih
i
?
和st s_ts
t
?
，前者是輸入序列第i個位置產(chǎn)生的隱狀態(tài)，后者是輸出序列在第t個位置產(chǎn)生的隱狀態(tài)。
所謂self-attention實際上就是，輸出序列就是輸入序列！因此，計算自己的attention得分，就叫做self-attention！
Context-attention是什么？
知道了self-attention，那你肯定猜到了context-attention是什么了：它是encoder和decoder之間的attention！所以，你也可以稱之為encoder-decoder attention!
context-attention一詞并不是本人原創(chuàng)，有些文章或者代碼會這樣描述，我覺得挺形象的，所以在此沿用這個稱呼。其他文章可能會有其他名稱，但是不要緊，我們抓住了重點即可，那就是兩個不同序列之間的attention，與self-attention相區(qū)別。
不管是self-attention還是context-attention，它們計算attention分數(shù)的時候，可以選擇很多方式，比如上面表中提到的：
additive attention
local-base
general
dot-product
scaled dot-product
那么我們的Transformer模型，采用的是哪種呢？答案是：scaled dot-product attention。
Scaled dot-product attention是什么？
論文Attention is all you need里面對于attention機制的描述是這樣的：
An attention function can be described as a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility of the query with the corresponding key.
這句話描述得很清楚了。翻譯過來就是：通過確定Q和K之間的相似程度來選擇V！
用公式來描述更加清晰：
Attention(Q,K,V)=softmax(QKTd√k)V \text{Attention}(Q,K,V)=softmax(\frac{QK^T}{\sqrt d_k})V
Attention(Q,K,V)=softmax(
d
?
k
?
QK
T
?
)V
scaled dot-product attention和dot-product attention唯一的區(qū)別就是，scaled dot-product attention有一個縮放因子
1d√k \frac{1}{\sqrt d_k}
d
?
k
?
1
?
。
上面公式中的
dk d_k
d
k
?
表示的是K的維度，在論文里面，默認是64。
那么為什么需要加上這個縮放因子呢？論文里給出了解釋：對于
dk d_k
d
k
?
很大的時候，點積得到的結(jié)果維度很大，使得結(jié)果處于softmax函數(shù)梯度很小的區(qū)域。
我們知道，梯度很小的情況，這對反向傳播不利。為了克服這個負面影響，除以一個縮放因子，可以一定程度上減緩這種情況。
為什么是
1d√k \frac{1}{\sqrt d_k}
d
?
k
?
1
?
呢？論文沒有進一步說明。個人覺得你可以使用其他縮放因子，看看模型效果有沒有提升。
論文也提供了一張很清晰的結(jié)構(gòu)圖，供大家參考：
Figure 3. Scaled dot-product attention architecture.
首先說明一下我們的K、Q、V是什么：
在encoder的self-attention中，Q、K、V都來自同一個地方（相等），他們是上一層encoder的輸出。對于第一層encoder，它們就是word embedding和positional encoding相加得到的輸入。
在decoder的self-attention中，Q、K、V都來自于同一個地方（相等），它們是上一層decoder的輸出。對于第一層decoder，它們就是word embedding和positional encoding相加得到的輸入。但是對于decoder，我們不希望它能獲得下一個time step（即將來的信息），因此我們需要進行sequence masking。
在encoder-decoder attention中，Q來自于decoder的上一層的輸出，K和V來自于encoder的輸出，K和V是一樣的。
Q、K、V三者的維度一樣，即
dq=dk=dv d_q=d_k=d_v
d
q
?
=d
k
?
=d
v
?
。
上面scaled dot-product attention和decoder的self-attention都出現(xiàn)了masking這樣一個東西。那么這個mask到底是什么呢？這兩處的mask操作是一樣的嗎？這個問題在后面會有詳細解釋。
Scaled dot-product attention的實現(xiàn)
咱們先把scaled dot-product attention實現(xiàn)了吧。代碼如下：
import torch
import torch.nn as nn
class ScaledDotProductAttention(nn.Module):
"""Scaled dot-product attention mechanism."""
def __init__(self, attention_dropout=0.0):
super(ScaledDotProductAttention, self).__init__()
self.dropout = nn.Dropout(attention_dropout)
self.softmax = nn.Softmax(dim=2)
def forward(self, q, k, v, scale=None, attn_mask=None):
"""前向傳播.
Args:
q: Queries張量，形狀為[B, L_q, D_q]
k: Keys張量，形狀為[B, L_k, D_k]
v: Values張量，形狀為[B, L_v, D_v]，一般來說就是k
scale: 縮放因子，一個浮點標量
attn_mask: Masking張量，形狀為[B, L_q, L_k]
Returns:
上下文張量和attetention張量
"""
attention = torch.bmm(q, k.transpose(1, 2))
if scale:
attention = attention * scale
if attn_mask:
# 給需要mask的地方設(shè)置一個負無窮
attention = attention.masked_fill_(attn_mask, -np.inf)
# 計算softmax
attention = self.softmax(attention)
# 添加dropout
attention = self.dropout(attention)
# 和V做點積
context = torch.bmm(attention, v)
return context, attention
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
Multi-head attention又是什么呢？
理解了Scaled dot-product attention，Multi-head attention也很簡單了。論文提到，他們發(fā)現(xiàn)將Q、K、V通過一個線性映射之后，分成
h h
h
份，對每一份進行scaled dot-product attention效果更好。然后，把各個部分的結(jié)果合并起來，再次經(jīng)過線性映射，得到最終的輸出。這就是所謂的multi-head attention。上面的超參數(shù)
h h
h
就是heads數(shù)量。論文默認是8。
下面是multi-head attention的結(jié)構(gòu)圖：
Figure 4: Multi-head attention architecture.
值得注意的是，上面所說的分成
h h
h
份是在
dk、dq、dv d_k、d_q、d_v
d
k
?
、d
q
?
、d
v
?
維度上面進行切分的。因此，進入到scaled dot-product attention的
dk d_k
d
k
?
實際上等于未進入之前的
DK/h D_K/h
D
K
?
/h
。
Multi-head attention允許模型加入不同位置的表示子空間的信息。
Multi-head attention的公式如下：
MultiHead(Q,K,V)=Concat(head1,…,headh)WO \text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_ 1,\dots,\text{head}_ h)W^O
MultiHead(Q,K,V)=Concat(head
1
?
,…,head
h
?
)W
O
其中，
headi=Attention(QWQi,KWKi,VWVi) \text{head}_ i = \text{Attention}(QW_i^Q,KW_i^K,VW_i^V)
head
i
?
=Attention(QW
i
Q
?
,KW
i
K
?
,VW
i
V
?
)
論文里面，
dmodel=512 d_{model}=512
d
model
?
=512
，
h=8 h=8
h=8
。所以在scaled dot-product attention里面的
dq=dk=dv=dmodel/h=512/8=64 d_q = d_k = d_v = d_{model}/h = 512/8 = 64
d
q
?
=d
k
?
=d
v
?
=d
model
?
/h=512/8=64
Multi-head attention的實現(xiàn)
相信大家已經(jīng)理清楚了multi-head attention，那么我們來實現(xiàn)它吧。代碼如下：
import torch
import torch.nn as nn
class MultiHeadAttention(nn.Module):
def __init__(self, model_dim=512, num_heads=8, dropout=0.0):
super(MultiHeadAttention, self).__init__()
self.dim_per_head = model_dim // num_heads
self.num_heads = num_heads
self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads)
self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads)
self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads)
self.dot_product_attention = ScaledDotProductAttention(dropout)
self.linear_final = nn.Linear(model_dim, model_dim)
self.dropout = nn.Dropout(dropout)
# multi-head attention之后需要做layer norm
self.layer_norm = nn.LayerNorm(model_dim)
def forward(self, key, value, query, attn_mask=None):
# 殘差連接
residual = query
dim_per_head = self.dim_per_head
num_heads = self.num_heads
batch_size = key.size(0)
# linear projection
key = self.linear_k(key)
value = self.linear_v(value)
query = self.linear_q(query)
# split by heads
key = key.view(batch_size * num_heads, -1, dim_per_head)
value = value.view(batch_size * num_heads, -1, dim_per_head)
query = query.view(batch_size * num_heads, -1, dim_per_head)
if attn_mask:
attn_mask = attn_mask.repeat(num_heads, 1, 1)
# scaled dot product attention
scale = (key.size(-1) // num_heads) ** -0.5
context, attention = self.dot_product_attention(
query, key, value, scale, attn_mask)
# concat heads
context = context.view(batch_size, -1, dim_per_head * num_heads)
# final linear projection
output = self.linear_final(context)
# dropout
output = self.dropout(output)
# add residual and norm layer
output = self.layer_norm(residual + output)
return output, attention
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
上面的代碼終于出現(xiàn)了Residual connection和Layer normalization。我們現(xiàn)在來解釋它們。
Residual connection是什么？
殘差連接其實很簡單！給你看一張示意圖你就明白了：
Figure 5. Residual connection.
假設(shè)網(wǎng)絡中某個層對輸入x作用后的輸出是
F(x) F(x)
F(x)
，那么增加residual connection之后，就變成了：
F(x)+x F(x)+x
F(x)+x
這個+x操作就是一個shortcut。
那么殘差結(jié)構(gòu)有什么好處呢？顯而易見：因為增加了一項x xx，那么該層網(wǎng)絡對x求偏導的時候，多了一個常數(shù)項
1 1
1
！所以在反向傳播過程中，梯度連乘，也不會造成梯度消失！
所以，代碼實現(xiàn)residual connection很非常簡單：
def residual(sublayer_fn,x):
return sublayer_fn(x)+x
1
2
文章開始的transformer架構(gòu)圖中的Add & Norm中的Add也就是指的這個shortcut。
至此，residual connection的問題理清楚了。更多關(guān)于殘差網(wǎng)絡的介紹可以看文末的參考文獻。
Layer normalization是什么？
GRADIENTS, BATCH NORMALIZATION AND LAYER NORMALIZATION一文對normalization有很好的解釋：
Normalization有很多種，但是它們都有一個共同的目的，那就是把輸入轉(zhuǎn)化成均值為0方差為1的數(shù)據(jù)。我們在把數(shù)據(jù)送入激活函數(shù)之前進行normalization（歸一化），因為我們不希望輸入數(shù)據(jù)落在激活函數(shù)的飽和區(qū)。
說到normalization，那就肯定得提到Batch Normalization。BN在CNN等地方用得很多。
BN的主要思想就是：在每一層的每一批數(shù)據(jù)上進行歸一化。
我們可能會對輸入數(shù)據(jù)進行歸一化，但是經(jīng)過該網(wǎng)絡層的作用后，我們的的數(shù)據(jù)已經(jīng)不再是歸一化的了。隨著這種情況的發(fā)展，數(shù)據(jù)的偏差越來越大，我的反向傳播需要考慮到這些大的偏差，這就迫使我們只能使用較小的學習率來防止梯度消失或者梯度爆炸。
BN的具體做法就是對每一小批數(shù)據(jù)，在批這個方向上做歸一化。如下圖所示：
Figure 6. Batch normalization example.(From theneuralperspective.com)
可以看到，右半邊求均值是沿著數(shù)據(jù)批量N的方向進行的！
Batch normalization的計算公式如下：
BN(xi)=α×xi?uBσ2B+?√+β BN(x_i)=\alpha\times\frac{x_i-u_B}{\sqrt{\sigma_B^2+\epsilon}}+\beta
BN(x
i
?
)=α×
σ
B
2
?
+?
?
x
i
?
?u
B
?
?
+β
具體的實現(xiàn)可以查看上圖的鏈接文章。
說完Batch normalization，就該說說咱們今天的主角Layer normalization。
那么什么是Layer normalization呢？:它也是歸一化數(shù)據(jù)的一種方式，不過LN是在每一個樣本上計算均值和方差，而不是BN那種在批方向計算均值和方差！
下面是LN的示意圖：
Figure 7. Layer normalization example.
和上面的BN示意圖一比較就可以看出二者的區(qū)別啦！
下面看一下LN的公式，也BN十分相似：
LN(xi)=α×xi?uLσ2L+?√+β LN(x_i)=\alpha\times\frac{x_i-u_L}{\sqrt{\sigma_L^2+\epsilon}}+\beta
LN(x
i
?
)=α×
σ
L
2
?
+?
?
x
i
?
?u
L
?
?
+β
Layer normalization的實現(xiàn)
上述兩個參數(shù)
α \alpha
α
和
β \beta
β
都是可學習參數(shù)。下面我們自己來實現(xiàn)Layer normalization(PyTorch已經(jīng)實現(xiàn)啦！)。代碼如下：
import torch
import torch.nn as nn
class LayerNorm(nn.Module):
"""實現(xiàn)LayerNorm。其實PyTorch已經(jīng)實現(xiàn)啦，見nn.LayerNorm。"""
def __init__(self, features, epsilon=1e-6):
"""Init.
Args:
features: 就是模型的維度。論文默認512
epsilon: 一個很小的數(shù)，防止數(shù)值計算的除0錯誤
"""
super(LayerNorm, self).__init__()
# alpha
self.gamma = nn.Parameter(torch.ones(features))
# beta
self.beta = nn.Parameter(torch.zeros(features))
self.epsilon = epsilon
def forward(self, x):
"""前向傳播.
Args:
x: 輸入序列張量，形狀為[B, L, D]
"""
# 根據(jù)公式進行歸一化
# 在X的最后一個維度求均值，最后一個維度就是模型的維度
mean = x.mean(-1, keepdim=True)
# 在X的最后一個維度求方差，最后一個維度就是模型的維度
std = x.std(-1, keepdim=True)
return self.gamma * (x - mean) / (std + self.epsilon) + self.beta
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
順便提一句，Layer normalization多用于RNN這種結(jié)構(gòu)。
Mask是什么？
現(xiàn)在終于輪到講解mask了!mask顧名思義就是掩碼，在我們這里的意思大概就是對某些值進行掩蓋，使其不產(chǎn)生效果。
需要說明的是，我們的Transformer模型里面涉及兩種mask。分別是padding mask和sequence mask。其中后者我們已經(jīng)在decoder的self-attention里面見過啦！
其中，padding mask在所有的scaled dot-product attention里面都需要用到，而sequence mask只有在decoder的self-attention里面用到。
所以，我們之前ScaledDotProductAttention的forward方法里面的參數(shù)attn_mask在不同的地方會有不同的含義。這一點我們會在后面說明。
Padding mask
什么是padding mask呢？回想一下，我們的每個批次輸入序列長度是不一樣的！也就是說，我們要對輸入序列進行對齊！具體來說，就是給在較短的序列后面填充0。因為這些填充的位置，其實是沒什么意義的，所以我們的attention機制不應該把注意力放在這些位置上，所以我們需要進行一些處理。
具體的做法是，把這些位置的值加上一個非常大的負數(shù)(可以是負無窮)，這樣的話，經(jīng)過softmax，這些位置的概率就會接近0！
而我們的padding mask實際上是一個張量，每個值都是一個Boolen，值為False的地方就是我們要進行處理的地方。
下面是實現(xiàn)：
def padding_mask(seq_k, seq_q):
# seq_k和seq_q的形狀都是[B,L]
len_q = seq_q.size(1)
# `PAD` is 0
pad_mask = seq_k.eq(0)
pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1) # shape [B, L_q, L_k]
return pad_mask
1
2
3
4
5
6
7
8
Sequence mask
文章前面也提到，sequence mask是為了使得decoder不能看見未來的信息。也就是對于一個序列，在time_step為t的時刻，我們的解碼輸出應該只能依賴于t時刻之前的輸出，而不能依賴t之后的輸出。因此我們需要想一個辦法，把t之后的信息給隱藏起來。
那么具體怎么做呢？也很簡單：產(chǎn)生一個上三角矩陣，上三角的值全為1，下三角的值權(quán)威0，對角線也是0。把這個矩陣作用在每一個序列上，就可以達到我們的目的啦。
具體的代碼實現(xiàn)如下：
def sequence_mask(seq):
batch_size, seq_len = seq.size()
mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8),
diagonal=1)
mask = mask.unsqueeze(0).expand(batch_size, -1, -1) # [B, L, L]
return mask
1
2
3
4
5
6
哈佛大學的文章The Annotated Transformer有一張效果圖:
Figure 8. Sequence mask.
值得注意的是，本來mask只需要二維的矩陣即可，但是考慮到我們的輸入序列都是批量的，所以我們要把原本二維的矩陣擴張成3維的張量。上面的代碼可以看出，我們已經(jīng)進行了處理。
回到本小結(jié)開始的問題，attn_mask參數(shù)有幾種情況？分別是什么意思？
對于decoder的self-attention，里面使用到的scaled dot-product attention，同時需要padding mask和sequence mask作為attn_mask，具體實現(xiàn)就是兩個mask相加作為attn_mask。
其他情況，attn_mask一律等于padding mask。
至此，mask相關(guān)的問題解決了。
Positional encoding是什么？
好了，終于要解釋位置編碼了，那就是文字開始的結(jié)構(gòu)圖提到的Positional encoding。
就目前而言，我們的Transformer架構(gòu)似乎少了點什么東西。沒錯，就是它對序列的順序沒有約束！我們知道序列的順序是一個很重要的信息，如果缺失了這個信息，可能我們的結(jié)果就是：所有詞語都對了，但是無法組成有意義的語句！
為了解決這個問題。論文提出了Positional encoding。這是啥？一句話概括就是：對序列中的詞語出現(xiàn)的位置進行編碼！如果對位置進行編碼，那么我們的模型就可以捕捉順序信息！
那么具體怎么做呢？論文的實現(xiàn)很有意思，使用正余弦函數(shù)。公式如下：
PE(pos,2i)=sin(pos/100002i/dmodel) PE(pos,2i) = sin(pos/10000^{2i/d_{model}})
PE(pos,2i)=sin(pos/10000
2i/d
model
?
)
PE(pos,2i+1)=cos(pos/100002i/dmodel) PE(pos,2i+1) = cos(pos/10000^{2i/d_{model}})
PE(pos,2i+1)=cos(pos/10000
2i/d
model
?
)
其中，pos是指詞語在序列中的位置。可以看出，在偶數(shù)位置，使用正弦編碼，在奇數(shù)位置，使用余弦編碼。
上面公式中的
dmodel d_{model}
d
model
?
是模型的維度，論文默認是512。
這個編碼公式的意思就是：給定詞語的位置
pos \text{pos}
pos
，我們可以把它編碼成
dmodel d_{model}
d
model
?
維的向量！也就是說，位置編碼的每一個維度對應正弦曲線，波長構(gòu)成了從
2π 2\pi
2π
到
10000?2π 10000*2\pi
10000?2π
的等比序列。
上面的位置編碼是絕對位置編碼。但是詞語的相對位置也非常重要。這就是論文為什么要使用三角函數(shù)的原因！
正弦函數(shù)能夠表達相對位置信息。，主要數(shù)學依據(jù)是以下兩個公式：
sin(α+β)=sinαcosβ+cosαsinβ sin(\alpha+\beta) = sin\alpha cos\beta + cos\alpha sin\beta
sin(α+β)=sinαcosβ+cosαsinβ
cos(α+β)=cosαcosβ?sinαsinβ cos(\alpha+\beta) = cos\alpha cos\beta - sin\alpha sin\beta
cos(α+β)=cosαcosβ?sinαsinβ
上面的公式說明，對于詞匯之間的位置偏移k，
PE(pos+k) PE(pos+k)
PE(pos+k)
可以表示成
PE(pos) PE(pos)
PE(pos)
和
PE(k) PE(k)
PE(k)
的組合形式，這就是表達相對位置的能力！
以上就是
PE PE
PE
的所有秘密。說完了positional encoding，那么我們還有一個與之處于同一地位的word embedding。
Word embedding大家都很熟悉了，它是對序列中的詞匯的編碼，把每一個詞匯編碼成
dmodel d_{model}
d
model
?
維的向量！看到?jīng)]有，Postional encoding是對詞匯的位置編碼，word embedding是對詞匯本身編碼！
所以，我更喜歡positional encoding的另外一個名字Positional embedding！
Positional encoding的實現(xiàn)
PE的實現(xiàn)也不難，按照論文的公式即可。代碼如下：
import torch
import torch.nn as nn
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_seq_len):
"""初始化。
Args:
d_model: 一個標量。模型的維度，論文默認是512
max_seq_len: 一個標量。文本序列的最大長度
"""
super(PositionalEncoding, self).__init__()
# 根據(jù)論文給的公式，構(gòu)造出PE矩陣
position_encoding = np.array([
[pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)]
for pos in range(max_seq_len)])
# 偶數(shù)列使用sin，奇數(shù)列使用cos
position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2])
position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2])
# 在PE矩陣的第一行，加上一行全是0的向量，代表這`PAD`的positional encoding
# 在word embedding中也經(jīng)常會加上`UNK`，代表位置單詞的word embedding，兩者十分類似
# 那么為什么需要這個額外的PAD的編碼呢？很簡單，因為文本序列的長度不一，我們需要對齊，
# 短的序列我們使用0在結(jié)尾補全，我們也需要這些補全位置的編碼，也就是`PAD`對應的位置編碼
pad_row = torch.zeros([1, d_model])
position_encoding = torch.cat((pad_row, position_encoding))
# 嵌入操作，+1是因為增加了`PAD`這個補全位置的編碼，
# Word embedding中如果詞典增加`UNK`，我們也需要+1。看吧，兩者十分相似
self.position_encoding = nn.Embedding(max_seq_len + 1, d_model)
self.position_encoding.weight = nn.Parameter(position_encoding,
requires_grad=False)
def forward(self, input_len):
"""神經(jīng)網(wǎng)絡的前向傳播。
Args:
input_len: 一個張量，形狀為[BATCH_SIZE, 1]。每一個張量的值代表這一批文本序列中對應的長度。
Returns:
返回這一批序列的位置編碼，進行了對齊。
"""
# 找出這一批序列的最大長度
max_len = torch.max(input_len)
tensor = torch.cuda.LongTensor if input_len.is_cuda else torch.LongTensor
# 對每一個序列的位置進行對齊，在原序列位置的后面補上0
# 這里range從1開始也是因為要避開PAD(0)的位置
input_pos = tensor(
[list(range(1, len + 1)) + [0] * (max_len - len) for len in input_len])
return self.position_encoding(input_pos)
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
Word embedding的實現(xiàn)
Word embedding應該是老生常談了，它實際上就是一個二維浮點矩陣，里面的權(quán)重是可訓練參數(shù)，我們只需要把這個矩陣構(gòu)建出來就完成了word embedding的工作。
所以，具體的實現(xiàn)很簡單：
import torch.nn as nn
embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0)
# 獲得輸入的詞嵌入編碼
seq_embedding = seq_embedding(inputs)*np.sqrt(d_model)
1
2
3
4
5
6
上面vocab_size就是詞典的大小，embedding_size就是詞嵌入的維度大小，論文里面就是等于
dmodel=512 d_{model}=512
d
model
?
=512
。所以word embedding矩陣就是一個vocab_size*embedding_size的二維張量。
如果你想獲取更詳細的關(guān)于word embedding的信息，可以看我的另外一個文章word2vec的筆記和實現(xiàn)。
Position-wise Feed-Forward network是什么？
這就是一個全連接網(wǎng)絡，包含兩個線性變換和一個非線性函數(shù)（實際上就是ReLU）。公式如下：
FFN(x)=max(0,xW1+b1)W2+b2 FFN(x)=max(0,xW_1+b_1)W_2+b_2
FFN(x)=max(0,xW
1
?
+b
1
?
)W
2
?
+b
2
?
這個線性變換在不同的位置都表現(xiàn)地一樣，并且在不同的層之間使用不同的參數(shù)。
論文提到，這個公式還可以用兩個核大小為1的一維卷積來解釋，卷積的輸入輸出都是
dmodel=512 d_{model}=512
d
model
?
=512
，中間層的維度是
dff=2048 d_{ff}=2048
d
ff
?
=2048
。
實現(xiàn)如下：
import torch
import torch.nn as nn
class PositionalWiseFeedForward(nn.Module):
def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0):
super(PositionalWiseFeedForward, self).__init__()
self.w1 = nn.Conv1d(model_dim, ffn_dim, 1)
self.w2 = nn.Conv1d(model_dim, ffn_dim, 1)
self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(model_dim)
def forward(self, x):
output = x.transpose(1, 2)
output = self.w2(F.relu(self.w1(output)))
output = self.dropout(output.transpose(1, 2))
# add residual and norm layer
output = self.layer_norm(x + output)
return output
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
Transformer的實現(xiàn)
至此，所有的細節(jié)都已經(jīng)解釋完了?，F(xiàn)在來完成我們Transformer模型的代碼。
首先，我們需要實現(xiàn)6層的encoder和decoder。
encoder代碼實現(xiàn)如下：
import torch
import torch.nn as nn
class EncoderLayer(nn.Module):
"""Encoder的一層。"""
def __init__(self, model_dim=512, num_heads=8, ffn_dim=2018, dropout=0.0):
super(EncoderLayer, self).__init__()
self.attention = MultiHeadAttention(model_dim, num_heads, dropout)
self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)
def forward(self, inputs, attn_mask=None):
# self attention
context, attention = self.attention(inputs, inputs, inputs, padding_mask)
# feed forward network
output = self.feed_forward(context)
return output, attention
class Encoder(nn.Module):
"""多層EncoderLayer組成Encoder。"""
def __init__(self,
vocab_size,
max_seq_len,
num_layers=6,
model_dim=512,
num_heads=8,
ffn_dim=2048,
dropout=0.0):
super(Encoder, self).__init__()
self.encoder_layers = nn.ModuleList(
[EncoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in
range(num_layers)])
self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)
self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)
def forward(self, inputs, inputs_len):
output = self.seq_embedding(inputs)
output += self.pos_embedding(inputs_len)
self_attention_mask = padding_mask(inputs, inputs)
attentions = []
for encoder in self.encoder_layers:
output, attention = encoder(output, self_attention_mask)
attentions.append(attention)
return output, attentions
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
通過文章前面的分析，代碼不需要更多解釋了。同樣的，我們的decoder代碼如下：
import torch
import torch.nn as nn
class DecoderLayer(nn.Module):
def __init__(self, model_dim, num_heads=8, ffn_dim=2048, dropout=0.0):
super(DecoderLayer, self).__init__()
self.attention = MultiHeadAttention(model_dim, num_heads, dropout)
self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)
def forward(self,
dec_inputs,
enc_outputs,
self_attn_mask=None,
context_attn_mask=None):
# self attention, all inputs are decoder inputs
dec_output, self_attention = self.attention(
dec_inputs, dec_inputs, dec_inputs, self_attn_mask)
# context attention
# query is decoder's outputs, key and value are encoder's inputs
dec_output, context_attention = self.attention(
enc_outputs, enc_outputs, dec_output, context_attn_mask)
# decoder's output, or context
dec_output = self.feed_forward(dec_output)
return dec_output, self_attention, context_attention
class Decoder(nn.Module):
def __init__(self,
vocab_size,
max_seq_len,
num_layers=6,
model_dim=512,
num_heads=8,
ffn_dim=2048,
dropout=0.0):
super(Decoder, self).__init__()
self.num_layers = num_layers
self.decoder_layers = nn.ModuleList(
[DecoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in
range(num_layers)])
self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)
self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)
def forward(self, inputs, inputs_len, enc_output, context_attn_mask=None):
output = self.seq_embedding(inputs)
output += self.pos_embedding(inputs_len)
self_attention_padding_mask = padding_mask(inputs, inputs)
seq_mask = sequence_mask(inputs)
self_attn_mask = torch.gt((self_attention_padding_mask + seq_mask), 0)
self_attentions = []
context_attentions = []
for decoder in self.decoder_layers:
output, self_attn, context_attn = decoder(
output, enc_output, self_attn_mask, context_attn_mask)
self_attentions.append(self_attn)
context_attentions.append(context_attn)
return output, self_attentions, context_attentions
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
最后，我們把encoder和decoder組成Transformer模型！
代碼如下：
import torch
import torch.nn as nn
class Transformer(nn.Module):
def __init__(self,
src_vocab_size,
src_max_len,
tgt_vocab_size,
tgt_max_len,
num_layers=6,
model_dim=512,
num_heads=8,
ffn_dim=2048,
dropout=0.2):
super(Transformer, self).__init__()
self.encoder = Encoder(src_vocab_size, src_max_len, num_layers, model_dim,
num_heads, ffn_dim, dropout)
self.decoder = Decoder(tgt_vocab_size, tgt_max_len, num_layers, model_dim,
num_heads, ffn_dim, dropout)
self.linear = nn.Linear(model_dim, tgt_vocab_size, bias=False)
self.softmax = nn.Softmax(dim=2)
def forward(self, src_seq, src_len, tgt_seq, tgt_len):
context_attn_mask = padding_mask(tgt_seq, src_seq)
output, enc_self_attn = self.encoder(src_seq, src_len)
output, dec_self_attn, ctx_attn = self.decoder(
tgt_seq, tgt_len, output, context_attn_mask)
output = self.linear(output)
output = self.softmax(output)
return output, enc_self_attn, dec_self_attn, ctx_attn
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
至此，Transformer模型已經(jīng)實現(xiàn)了！
參考文章
1.為什么ResNet和DenseNet可以這么深？一文詳解殘差塊為何有助于解決梯度彌散問題
2.GRADIENTS, BATCH NORMALIZATION AND LAYER NORMALIZATION
3.The Annotated Transformer
4.Building the Mighty Transformer for Sequence Tagging in PyTorch : Part I
5.Building the Mighty Transformer for Sequence Tagging in PyTorch : Part II
6.Attention?Attention!
參考代碼
1.jadore801120/attention-is-all-you-need-pytorch
2.JayParks/transformer
聯(lián)系我
Email: stupidme.me.lzy@gmail.com
WeChat: luozhouyang0528
---------------------
作者：luozhouyang
來源：CSDN
原文：https://blog.csdn.net/stupid_3/article/details/83184691
版權(quán)聲明：本文為博主原創(chuàng)文章，轉(zhuǎn)載請附上博文鏈接！
總結(jié)
以上是生活随笔為你收集整理的Transformer的PyTorch实现的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： python yield 和 yield
- 下一篇： Facebook 发布深度学习工具包 P