官網鏈接

Language Modeling with nn.Transformer and torchtext — PyTorch Tutorials 2.0.1+cu117 documentation

使用 NN.TRANSFORMER 和 TORCHTEXT進行語言建模

這是一個關于訓練模型使用nn.Transformer來預測序列中的下一個單詞的教程。

PyTorch 1.2版本包含了一個基于論文Attention is All You Need的標準transformer模塊。與循環(huán)神經網絡(RNNs)相比，transformer模型已被證明在許多序列對序列任務中具有更高的質量，同時具有更高的并發(fā)性。

nn.Transformer 模塊完全依賴于注意力機制(作為nn.MultiheadAttention實現(xiàn))來繪制輸入和輸出之間的全局依賴關系。nn.Transformer 模塊是高度模塊化的，這樣一個單一的組件(例如，nn.TransformerEncoder)可以很容易地使用/組合。

定義模型

在本教程中，我們在語言建模任務上訓練一個nn.TransformerEncoder模型。請注意，本教程不包括nn.TransformerDecoder的訓練，如上圖右半部分所示。語言建模任務是為給定單詞(或單詞序列)跟隨單詞序列的可能性分配一個概率。首先將一系列標記傳遞給源文本嵌入層，然后是一個位置編碼器來解釋單詞的順序(請參閱下一段了解更多細節(jié))。nn.TransformerEncoder 由nn.TransformerEncoderLayer的多個層組成。除了輸入序列外，還需要一個方形注意掩碼，因為在nn.TransformerDecoder中的自注意層只允許參與序列中較早的位置。對于語言建模任務，應掩蓋未來位置上的任何標記。為了生成輸出詞的概率分布，nn.TransformerEncoder 模型的輸出通過一個線性層來輸出非規(guī)范化的對數。這里沒有應用log-softmax函數，因為稍后會使用CrossEntropyLoss，它要求輸入是非標準化的對數。

import math
import os
from tempfile import TemporaryDirectory
from typing import Tupleimport torch
from torch import nn, Tensor
from torch.nn import TransformerEncoder, TransformerEncoderLayer
from torch.utils.data import datasetclass TransformerModel(nn.Module):def __init__(self, ntoken: int, d_model: int, nhead: int, d_hid: int,nlayers: int, dropout: float = 0.5):super().__init__()self.model_type = 'Transformer'self.pos_encoder = PositionalEncoding(d_model, dropout)encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout)self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)self.embedding = nn.Embedding(ntoken, d_model)self.d_model = d_modelself.linear = nn.Linear(d_model, ntoken)self.init_weights()def init_weights(self) -> None:initrange = 0.1self.embedding.weight.data.uniform_(-initrange, initrange)self.linear.bias.data.zero_()self.linear.weight.data.uniform_(-initrange, initrange)def forward(self, src: Tensor, src_mask: Tensor = None) -> Tensor:"""Arguments:src: Tensor, shape ``[seq_len, batch_size]``src_mask: Tensor, shape ``[seq_len, seq_len]``Returns:output Tensor of shape ``[seq_len, batch_size, ntoken]``"""src = self.embedding(src) * math.sqrt(self.d_model)src = self.pos_encoder(src)output = self.transformer_encoder(src, src_mask)output = self.linear(output)return output

PositionalEncoding模塊注入一些關于序列中記號的相對或絕對位置的信息。位置編碼器與文本嵌入層具有相同的維數，因此兩者可以相加。這里，我們使用不同頻率的正弦和余弦函數。

class PositionalEncoding(nn.Module):def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000):super().__init__()self.dropout = nn.Dropout(p=dropout)position = torch.arange(max_len).unsqueeze(1)div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))pe = torch.zeros(max_len, 1, d_model)pe[:, 0, 0::2] = torch.sin(position * div_term)pe[:, 0, 1::2] = torch.cos(position * div_term)self.register_buffer('pe', pe)def forward(self, x: Tensor) -> Tensor:"""Arguments:x: Tensor, shape ``[seq_len, batch_size, embedding_dim]``"""x = x + self.pe[:x.size(0)]return self.dropout(x)

加載和批處理數據

本教程使用torchtext生成Wikitext-2數據集。要訪問torchtext數據集，請按照GitHub - pytorch/data: A PyTorch repo for data loading and utilities to be shared by the PyTorch domain libraries. 的說明安裝torchdata.

%%bash
pip install portalocker
pip install torchdata

詞匯對象是基于訓練數據集構建的，用于將令牌數值化為張量。Wikitext-2將稀疏令牌表示為unk。給定順序數據的一維向量，batchify() 方法將數據排列到batch_size列中。如果數據沒有均勻地分成batch_size列，那么數據將被裁剪。例如，以字母表為數據(總長度為26)，batch_size=4，我們將字母表分成長度為6的序列，得到4個這樣的序列。

批處理支持更多的并行處理。然而，批處理意味著模型獨立處理每一列。在上面的例子中，G和F的依賴關系是無法學習的。

from torchtext.datasets import WikiText2
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iteratortrain_iter = WikiText2(split='train')
tokenizer = get_tokenizer('basic_english')
vocab = build_vocab_from_iterator(map(tokenizer, train_iter), specials=['<unk>'])
vocab.set_default_index(vocab['<unk>'])def data_process(raw_text_iter: dataset.IterableDataset) -> Tensor:"""Converts raw text into a flat Tensor."""data = [torch.tensor(vocab(tokenizer(item)), dtype=torch.long) for item in raw_text_iter]return torch.cat(tuple(filter(lambda t: t.numel() > 0, data)))# ``train_iter`` was "consumed" by the process of building the vocab,
# so we have to create it again
train_iter, val_iter, test_iter = WikiText2()
train_data = data_process(train_iter)
val_data = data_process(val_iter)
test_data = data_process(test_iter)device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')def batchify(data: Tensor, bsz: int) -> Tensor:"""Divides the data into ``bsz`` separate sequences, removing extra elementsthat wouldn't cleanly fit.Arguments:data: Tensor, shape ``[N]``bsz: int, batch sizeReturns:Tensor of shape ``[N // bsz, bsz]``"""seq_len = data.size(0) // bszdata = data[:seq_len * bsz]data = data.view(bsz, seq_len).t().contiguous()return data.to(device)batch_size = 20
eval_batch_size = 10
train_data = batchify(train_data, batch_size)  # shape ``[seq_len, batch_size]``
val_data = batchify(val_data, eval_batch_size)
test_data = batchify(test_data, eval_batch_size)

函數生成輸入和目標序列

get_batch() 為transformer模型生成一對輸入-目標序列。它將源數據細分為長度為bptt的塊。對于語言建模任務，模型需要以下單詞作為目標。例如，如果bptt值為2，那么當i = 0時，我們將得到以下兩個變量:

應該注意的是，數據塊的維度是0，與Transformer模型中的S維度一致。批次維度N 的維度是1.

bptt = 35
def get_batch(source: Tensor, i: int) -> Tuple[Tensor, Tensor]:"""Args:source: Tensor, shape ``[full_seq_len, batch_size]``i: intReturns:tuple (data, target), where data has shape ``[seq_len, batch_size]`` andtarget has shape ``[seq_len * batch_size]``"""seq_len = min(bptt, len(source) - 1 - i)data = source[i:i+seq_len]target = source[i+1:i+1+seq_len].reshape(-1)return data, target

初始化實例

模型超參數定義如下。詞匯表大小等于詞匯表對象的長度。

ntokens = len(vocab)  # size of vocabulary
emsize = 200  # embedding dimension
d_hid = 200  # dimension of the feedforward network model in ``nn.TransformerEncoder``
nlayers = 2  # number of ``nn.TransformerEncoderLayer`` in ``nn.TransformerEncoder``
nhead = 2  # number of heads in ``nn.MultiheadAttention``
dropout = 0.2  # dropout probability
model = TransformerModel(ntokens, emsize, nhead, d_hid, nlayers, dropout).to(device)

運行模型

我們將CrossEntropyLoss與 SGD (隨機梯度下降)優(yōu)化器一起使用。學習率最初設置為5.0，并遵循StepLR 計劃。在訓練期間，我們使用nn.utils.clip_grad_norm_ 防止梯度爆炸。

import timecriterion = nn.CrossEntropyLoss()
lr = 5.0  # learning rate
optimizer = torch.optim.SGD(model.parameters(), lr=lr)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)def train(model: nn.Module) -> None:model.train()  # turn on train modetotal_loss = 0.log_interval = 200start_time = time.time()num_batches = len(train_data) // bpttfor batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):data, targets = get_batch(train_data, i)output = model(data)output_flat = output.view(-1, ntokens)loss = criterion(output_flat, targets)optimizer.zero_grad()loss.backward()torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)optimizer.step()total_loss += loss.item()if batch % log_interval == 0 and batch > 0:lr = scheduler.get_last_lr()[0]ms_per_batch = (time.time() - start_time) * 1000 / log_intervalcur_loss = total_loss / log_intervalppl = math.exp(cur_loss)print(f'| epoch {epoch:3d} | {batch:5d}/{num_batches:5d} batches | 'f'lr {lr:02.2f} | ms/batch {ms_per_batch:5.2f} | 'f'loss {cur_loss:5.2f} | ppl {ppl:8.2f}')total_loss = 0start_time = time.time()def evaluate(model: nn.Module, eval_data: Tensor) -> float:model.eval()  # turn on evaluation modetotal_loss = 0.with torch.no_grad():for i in range(0, eval_data.size(0) - 1, bptt):data, targets = get_batch(eval_data, i)seq_len = data.size(0)output = model(data)output_flat = output.view(-1, ntokens)total_loss += seq_len * criterion(output_flat, targets).item()return total_loss / (len(eval_data) - 1)

循環(huán)epochs，如果驗證損失是目前為止我們看到的最好的，那么保存模型。在每個epoch之后調整學習速率。

best_val_loss = float('inf')
epochs = 3with TemporaryDirectory() as tempdir:best_model_params_path = os.path.join(tempdir, "best_model_params.pt")for epoch in range(1, epochs + 1):epoch_start_time = time.time()train(model)val_loss = evaluate(model, val_data)val_ppl = math.exp(val_loss)elapsed = time.time() - epoch_start_timeprint('-' * 89)print(f'| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | 'f'valid loss {val_loss:5.2f} | valid ppl {val_ppl:8.2f}')print('-' * 89)if val_loss < best_val_loss:best_val_loss = val_losstorch.save(model.state_dict(), best_model_params_path)scheduler.step()model.load_state_dict(torch.load(best_model_params_path)) # load best model states

輸出

| epoch   1 |   200/ 2928 batches | lr 5.00 | ms/batch 31.00 | loss  8.15 | ppl  3449.06
| epoch   1 |   400/ 2928 batches | lr 5.00 | ms/batch 28.73 | loss  6.25 | ppl   517.05
| epoch   1 |   600/ 2928 batches | lr 5.00 | ms/batch 28.56 | loss  5.61 | ppl   274.25
| epoch   1 |   800/ 2928 batches | lr 5.00 | ms/batch 28.42 | loss  5.31 | ppl   202.30
| epoch   1 |  1000/ 2928 batches | lr 5.00 | ms/batch 28.33 | loss  4.95 | ppl   140.81
| epoch   1 |  1200/ 2928 batches | lr 5.00 | ms/batch 28.28 | loss  4.55 | ppl    94.20
| epoch   1 |  1400/ 2928 batches | lr 5.00 | ms/batch 28.36 | loss  4.21 | ppl    67.25
| epoch   1 |  1600/ 2928 batches | lr 5.00 | ms/batch 28.45 | loss  3.99 | ppl    54.28
| epoch   1 |  1800/ 2928 batches | lr 5.00 | ms/batch 28.65 | loss  3.74 | ppl    41.89
| epoch   1 |  2000/ 2928 batches | lr 5.00 | ms/batch 28.56 | loss  3.66 | ppl    38.71
| epoch   1 |  2200/ 2928 batches | lr 5.00 | ms/batch 28.67 | loss  3.48 | ppl    32.44
| epoch   1 |  2400/ 2928 batches | lr 5.00 | ms/batch 28.74 | loss  3.49 | ppl    32.78
| epoch   1 |  2600/ 2928 batches | lr 5.00 | ms/batch 28.60 | loss  3.38 | ppl    29.50
| epoch   1 |  2800/ 2928 batches | lr 5.00 | ms/batch 28.46 | loss  3.29 | ppl    26.94
-----------------------------------------------------------------------------------------
| end of epoch   1 | time: 86.92s | valid loss  2.06 | valid ppl     7.88
-----------------------------------------------------------------------------------------
| epoch   2 |   200/ 2928 batches | lr 4.75 | ms/batch 28.88 | loss  3.10 | ppl    22.18
| epoch   2 |   400/ 2928 batches | lr 4.75 | ms/batch 28.50 | loss  3.02 | ppl    20.55
| epoch   2 |   600/ 2928 batches | lr 4.75 | ms/batch 28.65 | loss  2.86 | ppl    17.50
| epoch   2 |   800/ 2928 batches | lr 4.75 | ms/batch 28.68 | loss  2.85 | ppl    17.28
| epoch   2 |  1000/ 2928 batches | lr 4.75 | ms/batch 28.59 | loss  2.67 | ppl    14.43
| epoch   2 |  1200/ 2928 batches | lr 4.75 | ms/batch 28.55 | loss  2.68 | ppl    14.57
| epoch   2 |  1400/ 2928 batches | lr 4.75 | ms/batch 28.51 | loss  2.72 | ppl    15.13
| epoch   2 |  1600/ 2928 batches | lr 4.75 | ms/batch 28.44 | loss  2.69 | ppl    14.71
| epoch   2 |  1800/ 2928 batches | lr 4.75 | ms/batch 28.46 | loss  2.60 | ppl    13.51
| epoch   2 |  2000/ 2928 batches | lr 4.75 | ms/batch 28.51 | loss  2.61 | ppl    13.60
| epoch   2 |  2200/ 2928 batches | lr 4.75 | ms/batch 28.64 | loss  2.57 | ppl    13.04
| epoch   2 |  2400/ 2928 batches | lr 4.75 | ms/batch 28.67 | loss  2.57 | ppl    13.08
| epoch   2 |  2600/ 2928 batches | lr 4.75 | ms/batch 28.56 | loss  2.57 | ppl    13.05
| epoch   2 |  2800/ 2928 batches | lr 4.75 | ms/batch 28.61 | loss  2.55 | ppl    12.81
-----------------------------------------------------------------------------------------
| end of epoch   2 | time: 86.63s | valid loss  1.83 | valid ppl     6.24
-----------------------------------------------------------------------------------------
| epoch   3 |   200/ 2928 batches | lr 4.51 | ms/batch 28.71 | loss  2.43 | ppl    11.35
| epoch   3 |   400/ 2928 batches | lr 4.51 | ms/batch 28.82 | loss  2.37 | ppl    10.65
| epoch   3 |   600/ 2928 batches | lr 4.51 | ms/batch 28.67 | loss  2.27 | ppl     9.64
| epoch   3 |   800/ 2928 batches | lr 4.51 | ms/batch 28.74 | loss  2.29 | ppl     9.83
| epoch   3 |  1000/ 2928 batches | lr 4.51 | ms/batch 28.55 | loss  2.22 | ppl     9.22
| epoch   3 |  1200/ 2928 batches | lr 4.51 | ms/batch 28.73 | loss  2.25 | ppl     9.48
| epoch   3 |  1400/ 2928 batches | lr 4.51 | ms/batch 28.57 | loss  2.29 | ppl     9.89
| epoch   3 |  1600/ 2928 batches | lr 4.51 | ms/batch 28.73 | loss  2.36 | ppl    10.62
| epoch   3 |  1800/ 2928 batches | lr 4.51 | ms/batch 28.52 | loss  2.20 | ppl     9.07
| epoch   3 |  2000/ 2928 batches | lr 4.51 | ms/batch 28.61 | loss  2.26 | ppl     9.57
| epoch   3 |  2200/ 2928 batches | lr 4.51 | ms/batch 28.53 | loss  2.20 | ppl     9.03
| epoch   3 |  2400/ 2928 batches | lr 4.51 | ms/batch 28.45 | loss  2.23 | ppl     9.26
| epoch   3 |  2600/ 2928 batches | lr 4.51 | ms/batch 28.56 | loss  2.21 | ppl     9.13
| epoch   3 |  2800/ 2928 batches | lr 4.51 | ms/batch 28.54 | loss  2.31 | ppl    10.03
-----------------------------------------------------------------------------------------
| end of epoch   3 | time: 86.63s | valid loss  1.28 | valid ppl     3.60
-----------------------------------------------------------------------------------------

在測試數據集上評估最佳模型

test_loss = evaluate(model, test_data)
test_ppl = math.exp(test_loss)
print('=' * 89)
print(f'| End of training | test loss {test_loss:5.2f} | 'f'test ppl {test_ppl:8.2f}')
print('=' * 89)

輸出

=========================================================================================
| End of training | test loss  1.26 | test ppl     3.54
=========================================================================================