文章目錄

• 本文內容
• 將Transformer看成黑盒
• Transformer的推理過程
• Transformer的訓練過程
• Pytorch中的nn.Transformer
• • nn.Transformer簡介
  • nn.Transformer的構造參數詳解
  • Transformer的forward參數詳解
  • • src和tgt
    • src_mask、tgt_mask和memory_mask
    • key_padding_mask
  • nn.Transformer的使用
• 實戰：使用nn.Transformer實現一個簡單的Copy任務
• 參考資料：

本文內容

Transformer是個相對復雜的模型，可能有些人和我一樣，學了也不會用，或者感覺自己懂了，但又不懂。本文將Transformer看做一個黑盒，然后講解Pytorch中nn.Transformer的使用。

本文包含內容如下：

Transformer的訓練過程講解

Transformer的推理過程講解

Transformer的入參和出參講解

nn.Transformer的各個參數講解

nn.Transformer的mask機制詳解

實戰：使用nn.Transformer訓練一個copy任務。

你可以在該項目找到本文的源碼

開始之前，我們先導入要用到的包：

import math import randomimport torch import torch.nn as nn

將Transformer看成黑盒

這是一張經典的Transformer模型圖：

我們現在將其變成黑盒，將其蓋住：

我們現在再來看下Transformer的輸入和輸出：

這里是一個翻譯任務中transformer的輸入和輸出。transformer的輸入包含兩部分：

• inputs: 原句子對應的tokens，且是完整句子。一般0表示句子開始(<bos>)，1表示句子結束(<eos>)，2為填充(<pad>)。填充的目的是為了讓不同長度的句子變為同一個長度，這樣才可以組成一個batch。在代碼中，該變量一般取名src。
• outputs(shifted right)：上一個階段的輸出。雖然名字叫outputs，但是它是輸入。最開始為0（<bos>），然后本次預測出“我”后，下次調用Transformer的該輸入就變成<bos> 我。在代碼中，該變量一般取名tgt。

Transformer的輸出是一個概率分布。

Transformer的推理過程

這里先講Transformer的推理過程，因為這個簡單。其實通過上面的講解，你可能已經清楚了。上面是Transformer推理的第一步，緊接著第二步如圖：

Transformer的推理過程就是這樣一遍一遍調用Transformer，直到輸出<eos>或達到句子最大長度為止。

通常真正在實戰時，Transformer的Encoder部分只需要執行一遍就行了，這里為了簡單起見，就整體重新執行。

Transformer的訓練過程

在Transformer推理時，我們是一個詞一個詞的輸出，但在訓練時這樣做效率太低了，所以我們會將target一次性給到Transformer（當然，你也可以按照推理過程做），如圖所示：

從圖上可以看出，Transformer的訓練過程和推理過程主要有以下幾點異同：

源輸入src相同：對于Transformer的inputs部分(src參數)一樣，都是要被翻譯的句子。

目標輸入tgt不同：在Transformer推理時，tgt是從<bos>開始，然后每次加入上一次的輸出（第二次輸入為<bos> 我）。但在訓練時是一次將“完整”的結果給到Transformer，這樣其實和一個一個給結果上一致（可參考該篇的Mask Attention部分）。這里還有一個細節，就是tgt比src少了一位，src是7個token，而tgt是6個token。這是因為我們在最后一次推理時，只會傳入前n-1個token。舉個例子：假設我們要預測<bos> 我 愛 你 <eos>（這里忽略pad），我們最后一次的輸入tgt是<bos> 我 愛 你（沒有<eos>），因此我們的輸入tgt一定不會出現目標的最后一個token，所以一般tgt處理時會將目標句子刪掉最后一個token。

輸出數量變多：在訓練時，transformer會一次輸出多個概率分布。例如上圖，我就的等價于是tgt為<bos>時的輸出，愛就等價于tgt為<bos> 我時的輸出，依次類推。當然在訓練時，得到輸出概率分布后就可以計算loss了，并不需要將概率分布再轉成對應的文字。注意這里也有個細節，我們的輸出數量是6，對應到token就是我 愛 你 <eos> <pad> <pad>，這里少的是<bos>，因為<bos>不需要預測。計算loss時，我們也是要和的這幾個token進行計算，所以我們的label不包含<bos>。代碼中通常命名為tgt_y

當得到transformer的輸出后，我們就可以計算loss了，計算過程如圖：

Pytorch中的nn.Transformer

nn.Transformer簡介

在Pytorch中已經為我們實現了Transformer，我們可以直接拿來用，但nn.Transformer和我們上圖的還是有點區別，具體如圖：

Transformer并沒有實現Embedding和Positional Encoding和最后的Linear+Softmax部分，這里我簡單對這幾部分進行說明：

• Embedding: 負責將token映射成高維向量。例如將123映射成[0.34, 0.45, 0.123, ..., 0.33]。通常使用nn.Embedding來實現。但nn.Embedding的參數并不是一成不變的，也是會參與梯度下降。關于nn.Embedding可參考文章Pytorch nn.Embedding的基本使用
• Positional Encoding：位置編碼。用于為token編碼增加位置信息，例如I love you這三個token編碼后的向量并不包含其位置信息(love左邊是I，右邊是you這個信息)。這個位置信息還挺重要的，有和沒有真的是天差地別。
• Linear+Softmax：一個線性層加一個Softmax，用于對nn.Transformer輸出的結果進行token預測。如果把Transformer比作CNN，那么nn.Transformer實現的就是卷積層，而Linear+Softmax就是卷積層后面的線性層。

這里我先簡單的演示一下nn.Transformer的使用：

# 定義編碼器，詞典大小為10，要把token編碼成128維的向量 embedding = nn.Embedding(10, 128) # 定義transformer，模型維度為128（也就是詞向量的維度） transformer = nn.Transformer(d_model=128, batch_first=True) # batch_first一定不要忘記 # 定義源句子，可以想想成是 <bos> 我 愛 吃 肉 和 菜 <eos> <pad> <pad> src = torch.LongTensor([[0, 3, 4, 5, 6, 7, 8, 1, 2, 2]]) # 定義目標句子，可以想想是 <bos> I like eat meat and vegetables <eos> <pad> tgt = torch.LongTensor([[0, 3, 4, 5, 6, 7, 8, 1, 2]]) # 將token編碼后送給transformer（這里暫時不加Positional Encoding） outputs = transformer(embedding(src), embedding(tgt)) outputs.size() torch.Size([1, 9, 128])

Transformer輸出的Shape和tgt編碼后的Shape一致。在訓練時，我們會把transformer的所有輸出送給Linear，而在推理時，只需要將最后一個輸出送給Linear即可，即outputs[:, -1]。

nn.Transformer的構造參數詳解

Transformer構造參數眾多，所以我們還需要將黑盒稍微打開一下：

nn.Transformer主要由兩部分構成：nn.TransformerEncoder和nn.TransformerDecoder。而nn.TransformerEncoder又是由多個nn.TransformerEncoderLayer堆疊而成的，圖中的Nx就是要堆疊多少層。nn.TransformerDecoder同理。

下面是nn.Transformer的構造參數：

• d_model: Encoder和Decoder輸入參數的特征維度。也就是詞向量的維度。默認為512
• nhead: 多頭注意力機制中，head的數量。關于Attention機制，可以參考這篇文章。注意該值并不影響網絡的深度和參數數量。默認值為8。
• num_encoder_layers: TransformerEncoderLayer的數量。該值越大，網絡越深，網絡參數量越多，計算量越大。默認值為6
• num_decoder_layers：TransformerDecoderLayer的數量。該值越大，網絡越深，網絡參數量越多，計算量越大。默認值為6
• dim_feedforward：Feed Forward層（Attention后面的全連接網絡）的隱藏層的神經元數量。該值越大，網絡參數量越多，計算量越大。默認值為2048
• dropout：dropout值。默認值為0.1
• activation： Feed Forward層的激活函數。取值可以是string(“relu” or “gelu”)或者一個一元可調用的函數。默認值是relu
• custom_encoder：自定義Encoder。若你不想用官方實現的TransformerEncoder，你可以自己實現一個。默認值為None
• custom_decoder: 自定義Decoder。若你不想用官方實現的TransformerDecoder，你可以自己實現一個。
• layer_norm_eps: Add&Norm層中，BatchNorm的eps參數值。默認為1e-5
• batch_first：batch維度是否是第一個。如果為True，則輸入的shape應為(batch_size, 詞數，詞向量維度)，否則應為(詞數, batch_size, 詞向量維度)。默認為False。這個要特別注意，因為大部分人的習慣都是將batch_size放在最前面，而這個參數的默認值又是False，所以會報錯。
• norm_first – 是否要先執行norm。例如，在圖中的執行順序為 Attention -> Add -> Norm。若該值為True，則執行順序變為：Norm -> Attention -> Add。

Transformer的forward參數詳解

Transformer的forward參數需要詳細解釋，這里我先將其列出來，進行粗略解釋，然后再逐個進行詳細解釋：

• src: Encoder的輸入。也就是將token進行Embedding并Positional Encoding之后的tensor。必填參數。Shape為(batch_size, 詞數, 詞向量維度)
• tgt: 與src同理，Decoder的輸入。 必填參數。Shape為(詞數, 詞向量維度)
• src_mask: 對src進行mask。不常用。Shape為(詞數, 詞數)
• tgt_mask：對tgt進行mask。常用。Shape為(詞數, 詞數)
• memory_mask – 對Encoder的輸出memory進行mask。 不常用。Shape為(batch_size, 詞數, 詞數)
• src_key_padding_mask：對src的token進行mask. 常用。Shape為(batch_size, 詞數)
• tgt_key_padding_mask：對tgt的token進行mask。常用。Shape為(batch_size, 詞數)
• memory_key_padding_mask：對tgt的token進行mask。不常用。Shape為(batch_size, 詞數)

上面的所有mask都是0代表不遮掩，-inf代表遮掩。另外，src_mask、tgt_mask和memory_mask是不需要傳batch的。

補充：上面的說法是pytorch1.11版本。我發現1.11版本key_padding_mask可以用True/False，但1.12版本key_padding_mask好像只能是True/False，其中True表示遮掩，而False表示不遮掩（這個可不要弄混了，這個和Transformer原論文實現正好相反，如果弄反了，會造成結果為nan）。

上面說了和沒說其實差不多，重要的是每個參數的是否常用和其對應的Shape（這里我默認batch_first=True）。 接下來對各個參數進行詳細解釋。

src和tgt

src參數和tgt參數分別為Encoder和Decoder的輸入參數。它們是對token進行編碼后，再經過Positional Encoding之后的結果。

例如：我們一開始的輸入為：[[0, 3, 4, 5, 6, 7, 8, 1, 2, 2]]，Shape為(1, 10)，表示batch_size為1, 每句10個詞。

在經過Embedding后，Shape就變成了(1, 10, 128)，表示batch_size為1, 每句10個詞，每個詞被編碼為了128維的向量。

src就是這個(1, 10, 128)的向量。tgt同理

src_mask、tgt_mask和memory_mask

要真正理解mask，需要學習Attention機制，可參考該篇。這里只做一個簡要的說明。

在經過Attention層時，會讓每個詞具有上下文關系，也就是每個詞除了自己的信息外，還包含其他詞的信息。例如：蘋果 很 好吃和蘋果 手機 很 好玩，這兩個蘋果顯然指的不是同一個意思。但讓蘋果這個詞具備了后面好吃或手機這兩個詞的信息后，那就可以區分這兩個蘋果的含義了。

在Attention中，我們有這么一個“方陣”，描述著詞與詞之間的關系，例如：

蘋果 很 好吃 蘋果 [[0.5, 0.1, 0.4], 很 [0.1, 0.8, 0.1], 好吃 [0.3, 0.1, 0.6],]

在上述矩陣中，蘋果這個詞與自身, 很和好吃三個詞的關系權重就是[0.5, 0.1, 0.4]，通過該矩陣，我們就可以得到包含上下文的蘋果了，即

$${\text{蘋果}}^{\prime }=\text{蘋果}\times 0.5+\text{很}\times 0.1+\text{好吃}\times 0.4$$

但在實際推理時，詞是一個一個輸出的。若蘋果很好吃是tgt的話，那么蘋果是不應該包含很和好吃的上下文信息的，所以我們希望為：

$${\text{蘋果}}^{\prime }=\text{蘋果}\times 0.5$$

同理，很字可以包含蘋果的上下信息，但不能包含好吃，所以為：

$${\text{很}}^{\prime }=\text{蘋果}\times 0.1+\text{很}\times 0.8$$

那要完成這個事情，那只需要改變方陣即可：

蘋果 很 好吃 蘋果 [[0.5, 0, 0], 很 [0.1, 0.8, 0], 好吃 [0.3, 0.1, 0.6],]

而這個事情我們就可以使用mask掩碼來完成，即：

蘋果 很 好吃 蘋果 [[ 0, -inf, -inf], 很 [ 0, 0, -inf], 好吃 [ 0, 0, 0]]

其中0表示不遮掩，而-inf表示遮掩。（之所以這么定是因為這個方陣還要過softmax，所以會使-inf變為0）。

所以，對于tgt_mask，我們只需要生成一個斜著覆蓋的方陣即可，我們可以利用nn.Transformer.generate_square_subsequent_mask來完成，例如：

nn.Transformer.generate_square_subsequent_mask(5) # 這個5指的是tgt的token的數量 tensor([[0., -inf, -inf, -inf, -inf],[0., 0., -inf, -inf, -inf],[0., 0., 0., -inf, -inf],[0., 0., 0., 0., -inf],[0., 0., 0., 0., 0.]])

通過上面的分析，src和memory一般是不需要進行mask的，所以不常用。

key_padding_mask

在我們的src和tgt語句中，除了本身的詞外，還包含了三種token: <bos>, <eos> 和 <pad>。這里面的<pad>只是為了改變句子長度，方便將不同長度的句子組成batch而進行填充的。該token沒有任何意義，所以在計算Attention時，也不想讓它們參與，所以也要mask。而對于這種mask就需要用到key_padding_mask這個參數了。

例如，我們的src為[[0, 3, 4, 5, 6, 7, 8, 1, 2, 2]]，其中2是<pad>，所以我們的src_key_padding_mask就應為[[0, 0, 0, 0, 0, 0, 0, 0, -inf, -inf]]，即將最后兩個2給掩蓋住。

tgt_key_padding_mask同理。但memory_key_padding_mask就沒有必要用了。

在Transformer的源碼中或其他實現中，tgt_mask和tgt_key_padding_mask是合在一起的，例如：

[[0., -inf, -inf, -inf], # tgt_mask[0., 0., -inf, -inf],[0., 0., 0., -inf],[0., 0., 0., 0.]]+[[0., 0., 0., -inf]] # tgt_key_padding_mask= [[0., -inf, -inf, -inf], # 合并之后的[0., 0., -inf, -inf],[0., 0., 0., -inf],[0., 0., 0., -inf]]

nn.Transformer的使用

接下來我們來簡單的使用一下nn.Transformer:

首先我們定義src和tgt:

src = torch.LongTensor([[0, 8, 3, 5, 5, 9, 6, 1, 2, 2, 2],[0, 6, 6, 8, 9, 1 ,2, 2, 2, 2, 2], ]) tgt = torch.LongTensor([[0, 8, 3, 5, 5, 9, 6, 1, 2, 2],[0, 6, 6, 8, 9, 1 ,2, 2, 2, 2], ])

接下來定義一個輔助函數來生成src_key_padding_mask和tgt_key_padding_mask:

def get_key_padding_mask(tokens):key_padding_mask = torch.zeros(tokens.size())key_padding_mask[tokens == 2] = -torch.infreturn key_padding_masksrc_key_padding_mask = get_key_padding_mask(src) tgt_key_padding_mask = get_key_padding_mask(tgt) print(tgt_key_padding_mask) tensor([[0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf],[0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf]])

然后通過Transformer內容方法生成tgt_mask：

tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size(-1)) print(tgt_mask) tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],[0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],[0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],[0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf],[0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf],[0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf],[0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf],[0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf],[0., 0., 0., 0., 0., 0., 0., 0., 0., -inf],[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])

之后就可以定義Embedding和Transformer進行調用了：

# 定義編碼器，詞典大小為10，要把token編碼成128維的向量 embedding = nn.Embedding(10, 128) # 定義transformer，模型維度為128（也就是詞向量的維度） transformer = nn.Transformer(d_model=128, batch_first=True) # batch_first一定不要忘記 # 將token編碼后送給transformer（這里暫時不加Positional Encoding） outputs = transformer(embedding(src), embedding(tgt),tgt_mask=tgt_mask,src_key_padding_mask=src_key_padding_mask,tgt_key_padding_mask=tgt_key_padding_mask) print(outputs.size()) torch.Size([2, 10, 128])

實戰：使用nn.Transformer實現一個簡單的Copy任務

任務描述：讓Transformer預測輸入。例如，輸入為[0, 3, 4, 6, 7, 1, 2, 2]，則期望的輸出為[0, 3, 4, 6, 7, 1]。

首先，我們定義一下句子的最大長度：

max_length=16

定義PositionEncoding類，不需要知道具體什么意思，直接拿過來用即可。

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)# 初始化Shape為(max_len, d_model)的PE (positional encoding)pe = torch.zeros(max_len, d_model)# 初始化一個tensor [[0, 1, 2, 3, ...]]position = torch.arange(0, max_len).unsqueeze(1)# 這里就是sin和cos括號中的內容，通過e和ln進行了變換div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))# 計算PE(pos, 2i)pe[:, 0::2] = torch.sin(position * div_term)# 計算PE(pos, 2i+1)pe[:, 1::2] = torch.cos(position * div_term)# 為了方便計算，在最外面在unsqueeze出一個batchpe = pe.unsqueeze(0)# 如果一個參數不參與梯度下降，但又希望保存model的時候將其保存下來# 這個時候就可以用register_bufferself.register_buffer("pe", pe)def forward(self, x):"""x 為embedding后的inputs，例如(1,7, 128)，batch size為1,7個單詞，單詞維度為128"""# 將x和positional encoding相加。x = x + self.pe[:, : x.size(1)].requires_grad_(False)return self.dropout(x)

定義我們的Copy模型：

class CopyTaskModel(nn.Module):def __init__(self, d_model=128):super(CopyTaskModel, self).__init__()# 定義詞向量，詞典數為10。我們不預測兩位小數。self.embedding = nn.Embedding(num_embeddings=10, embedding_dim=128)# 定義Transformer。超參是我拍腦袋想的self.transformer = nn.Transformer(d_model=128, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=512, batch_first=True)# 定義位置編碼器self.positional_encoding = PositionalEncoding(d_model, dropout=0)# 定義最后的線性層，這里并沒有用Softmax，因為沒必要。# 因為后面的CrossEntropyLoss中自帶了self.predictor = nn.Linear(128, 10)def forward(self, src, tgt):# 生成masktgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size()[-1])src_key_padding_mask = CopyTaskModel.get_key_padding_mask(src)tgt_key_padding_mask = CopyTaskModel.get_key_padding_mask(tgt)# 對src和tgt進行編碼src = self.embedding(src)tgt = self.embedding(tgt)# 給src和tgt的token增加位置信息src = self.positional_encoding(src)tgt = self.positional_encoding(tgt)# 將準備好的數據送給transformerout = self.transformer(src, tgt,tgt_mask=tgt_mask,src_key_padding_mask=src_key_padding_mask,tgt_key_padding_mask=tgt_key_padding_mask)"""這里直接返回transformer的結果。因為訓練和推理時的行為不一樣，所以在該模型外再進行線性層的預測。"""return out@staticmethoddef get_key_padding_mask(tokens):"""用于key_padding_mask"""key_padding_mask = torch.zeros(tokens.size())key_padding_mask[tokens == 2] = -torch.infreturn key_padding_mask model = CopyTaskModel()

這里簡單的嘗試下我們定義的模型：

src = torch.LongTensor([[0, 3, 4, 5, 6, 1, 2, 2]]) tgt = torch.LongTensor([[3, 4, 5, 6, 1, 2, 2]]) out = model(src, tgt) print(out.size()) print(out) torch.Size([1, 7, 128]) tensor([[[ 2.1870e-01, 1.3451e-01, 7.4523e-01, -1.1865e+00, -9.1054e-01,6.0285e-01, 8.3666e-02, 5.3425e-01, 2.2247e-01, -3.6559e-01,.... -9.1266e-01, 1.7342e-01, -5.7250e-02, 7.1583e-02, 7.0782e-01,-3.5137e-01, 5.1000e-01, -4.7047e-01]]],grad_fn=<NativeLayerNormBackward0>)

沒什么問題，那就接著定義損失函數和優化器，因為是多分類問題，所以用CrossEntropyLoss:

criteria = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)

接著再定義一個生成隨時數據的工具函數：

def generate_random_batch(batch_size, max_length=16):src = []for i in range(batch_size):# 隨機生成句子長度random_len = random.randint(1, max_length - 2)# 隨機生成句子詞匯，并在開頭和結尾增加<bos>和<eos>random_nums = [0] + [random.randint(3, 9) for _ in range(random_len)] + [1]# 如果句子長度不足max_length，進行填充random_nums = random_nums + [2] * (max_length - random_len - 2)src.append(random_nums)src = torch.LongTensor(src)# tgt不要最后一個tokentgt = src[:, :-1]# tgt_y不要第一個的tokentgt_y = src[:, 1:]# 計算tgt_y，即要預測的有效token的數量n_tokens = (tgt_y != 2).sum()# 這里的n_tokens指的是我們要預測的tgt_y中有多少有效的token，后面計算loss要用return src, tgt, tgt_y, n_tokens generate_random_batch(batch_size=2, max_length=6) (tensor([[0, 7, 6, 8, 7, 1],[0, 9, 4, 1, 2, 2]]),tensor([[0, 7, 6, 8, 7],[0, 9, 4, 1, 2]]),tensor([[7, 6, 8, 7, 1],[9, 4, 1, 2, 2]]),tensor(8))

開始進行訓練：

total_loss = 0for step in range(2000):# 生成數據src, tgt, tgt_y, n_tokens = generate_random_batch(batch_size=2, max_length=max_length)# 清空梯度optimizer.zero_grad()# 進行transformer的計算out = model(src, tgt)# 將結果送給最后的線性層進行預測out = model.predictor(out)"""計算損失。由于訓練時我們的是對所有的輸出都進行預測，所以需要對out進行reshape一下。我們的out的Shape為(batch_size, 詞數, 詞典大小)，view之后變為：(batch_size*詞數, 詞典大小)。而在這些預測結果中，我們只需要對非<pad>部分進行，所以需要進行正則化。也就是除以n_tokens。"""loss = criteria(out.contiguous().view(-1, out.size(-1)), tgt_y.contiguous().view(-1)) / n_tokens# 計算梯度loss.backward()# 更新參數optimizer.step()total_loss += loss# 每40次打印一下lossif step != 0 and step % 40 == 0:print("Step {}, total_loss: {}".format(step, total_loss))total_loss = 0 Step 40, total_loss: 3.570814609527588 Step 80, total_loss: 2.4842987060546875 ...略 Step 1920, total_loss: 0.4518987536430359 Step 1960, total_loss: 0.37290623784065247

在完成模型訓練后，我們來使用一下我們的模型：

model = model.eval() # 隨便定義一個src src = torch.LongTensor([[0, 4, 3, 4, 6, 8, 9, 9, 8, 1, 2, 2]]) # tgt從<bos>開始，看看能不能重新輸出src中的值 tgt = torch.LongTensor([[0]]) # 一個一個詞預測，直到預測為<eos>，或者達到句子最大長度 for i in range(max_length):# 進行transformer計算out = model(src, tgt)# 預測結果，因為只需要看最后一個詞，所以取`out[:, -1]`predict = model.predictor(out[:, -1])# 找出最大值的indexy = torch.argmax(predict, dim=1)# 和之前的預測結果拼接到一起tgt = torch.concat([tgt, y.unsqueeze(0)], dim=1)# 如果為<eos>，說明預測結束，跳出循環if y == 1:break print(tgt) tensor([[0, 4, 3, 4, 6, 8, 9, 9, 8, 1]])

可以看到，我們的模型成功預測了src的輸入。

---

參考資料：

nn.Transformer官方文檔: https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html

層層剖析，讓你徹底搞懂Self-Attention、MultiHead-Attention和Masked-Attention的機制和原理: https://blog.csdn.net/zhaohongfei_358/article/details/122861751

總結

以上是生活随笔為你收集整理的Pytorch中 nn.Transformer的使用详解与Transformer的黑盒讲解的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。