基于TorchText的语言翻译

译者：PengboLiu

校验：PengboLiu

本教程介绍如何使用torchtext的几个类来预处理英德数据集，该数据集可以用来训练seq2seq模型，既而能自动把德语句子翻译成英语。

本文基于PyTorch社区成员Ben Trevett的教程，并得到了他本人的许可。

阅读完本教程，你将能够：

Field和 TranslationDataset

torchtext具有创建数据集的功能，可以轻松对其迭代以构建机器翻译模型。一个关键的类是Filed，它指定每个句子的预处理方法，另一个类是TranslationDataset ; torchtext内置了几个翻译数据集；在本教程中，我们将使用 Multi30k dataset数据集，其中包含约30000个英德句对(平均长度约13个字）。

注：本教程中的tokenization 需要使用 Spacy 。Spacy包可以帮助我们对英语以外的语言tokenization。torchtext提供了basic_english的tokenizer ，但是对于其他语言，使用Spacy对我们而言是最好的选择。

为了运行该教程，首先要使用pip或conda安装Spacy。接下来，下载英德原始数据：

  1. python -m spacy download en
  2. python -m spacy download de

安装Spacy后，以下代码将根据Field中定义的tokenizer 处理TranslationDataset中的每个句子。

  1.  from torchtext.datasets import Multi30k
  2.  from torchtext.data import Field, BucketIterator
  3. SRC = Field(tokenize = "spacy",
  4.             tokenizer_language="de",
  5.             init_token = '<sos>',
  6.             eos_token = '<eos>',
  7.             lower = True)
  9. TRG = Field(tokenize = "spacy",
  10.             tokenizer_language="en",
  11.             init_token = '<sos>',
  12.             eos_token = '<eos>',
  13.             lower = True)
  15. train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'),
  16.                                                     fields = (SRC, TRG))

输出：
downloading training.tar.gz
downloading validation.tar.gz
downloading mmt_task1_test2016.tar.gz

现在，我们已经定义好了train_data，torchtext的Field有一个非常有用的功能 ：我们可以使用build_vocab方法创建每个语言的词汇表。

  1. SRC.build_vocab(train_data, min_freq = 2)
  2. TRG.build_vocab(train_data, min_freq = 2)

一旦这些代码行被运行，SRC.vocab.stoi将成为一个tokens作为key，索引作为value的词典；对应的， SRC.vocab.itos是一个交换了key和value内容相同的字典。在本教程中我们不会广泛使用此功能，但是你可能在遇到其他NLP任务有用。

BucketIterator

我们使用最后一个torchtext的特性是BucketIterator， 它以TranslationDataset作为第一个参数，所以易于使用。如文档所说：定义一个迭代器，该迭代器将相似长度的数据放在一起。产生每个新bacth时，最大程度地减少所需的填充量。

  1. import torch
  3. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  5. BATCH_SIZE = 128
  7. train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
  8.     (train_data, valid_data, test_data),
  9.     batch_size = BATCH_SIZE,
  10.     device = device)

可以像DataLoader一样调用这些迭代器。 在下面的训练和评估函数中，它们可以简单地通过以下方式调用：

  1. for i, batch in enumerate(iterator):

每个batch于是具有SRC和TRG属性：

  1. src = batch.src
  2. trg = batch.trg

定义我们的nn.Module和Optimizer

解决了数据集的问题并为之定义好迭代器，我们剩下的任务就是定义模型和优化器完成训练过程。

具体来说，我们的模型遵循此处描述的结构。

注意：我们选择这种模型并不是因为它是目前最优的，而是因为它是机器翻译的标准模型。众所周知，目前机器翻译的最优模型是Transformers。

  1. import random
  2. from typing import Tuple
  4. import torch.nn as nn
  5. import torch.optim as optim
  6. import torch.nn.functional as F
  7. from torch import Tensor
  10. class Encoder(nn.Module):
  11.     def __init__(self,
  12.                  input_dim: int,
  13.                  emb_dim: int,
  14.                  enc_hid_dim: int,
  15.                  dec_hid_dim: int,
  16.                  dropout: float):
  17.         super().__init__()
  19.         self.input_dim = input_dim
  20.         self.emb_dim = emb_dim
  21.         self.enc_hid_dim = enc_hid_dim
  22.         self.dec_hid_dim = dec_hid_dim
  23.         self.dropout = dropout
  25.         self.embedding = nn.Embedding(input_dim, emb_dim)
  27.         self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
  29.         self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
  31.         self.dropout = nn.Dropout(dropout)
  33.     def forward(self,
  34.                 src: Tensor) -> Tuple[Tensor]:
  36.         embedded = self.dropout(self.embedding(src))
  38.         outputs, hidden = self.rnn(embedded)
  40.         hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
  42.         return outputs, hidden
  45. class Attention(nn.Module):
  46.     def __init__(self,
  47.                  enc_hid_dim: int,
  48.                  dec_hid_dim: int,
  49.                  attn_dim: int):
  50.         super().__init__()
  52.         self.enc_hid_dim = enc_hid_dim
  53.         self.dec_hid_dim = dec_hid_dim
  55.         self.attn_in = (enc_hid_dim * 2) + dec_hid_dim
  57.         self.attn = nn.Linear(self.attn_in, attn_dim)
  59.     def forward(self,
  60.                 decoder_hidden: Tensor,
  61.                 encoder_outputs: Tensor) -> Tensor:
  63.         src_len = encoder_outputs.shape[0]
  65.         repeated_decoder_hidden = decoder_hidden.unsqueeze(1).repeat(1, src_len, 1)
  67.         encoder_outputs = encoder_outputs.permute(1, 0, 2)
  69.         energy = torch.tanh(self.attn(torch.cat((
  70.             repeated_decoder_hidden,
  71.             encoder_outputs),
  72.             dim = 2)))
  74.         attention = torch.sum(energy, dim=2)
  76.         return F.softmax(attention, dim=1)
  79. class Decoder(nn.Module):
  80.     def __init__(self,
  81.                  output_dim: int,
  82.                  emb_dim: int,
  83.                  enc_hid_dim: int,
  84.                  dec_hid_dim: int,
  85.                  dropout: int,
  86.                  attention: nn.Module):
  87.         super().__init__()
  89.         self.emb_dim = emb_dim
  90.         self.enc_hid_dim = enc_hid_dim
  91.         self.dec_hid_dim = dec_hid_dim
  92.         self.output_dim = output_dim
  93.         self.dropout = dropout
  94.         self.attention = attention
  96.         self.embedding = nn.Embedding(output_dim, emb_dim)
  98.         self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
  100.         self.out = nn.Linear(self.attention.attn_in + emb_dim, output_dim)
  102.         self.dropout = nn.Dropout(dropout)
  105.     def _weighted_encoder_rep(self,
  106.                               decoder_hidden: Tensor,
  107.                               encoder_outputs: Tensor) -> Tensor:
  109.         a = self.attention(decoder_hidden, encoder_outputs)
  111.         a = a.unsqueeze(1)
  113.         encoder_outputs = encoder_outputs.permute(1, 0, 2)
  115.         weighted_encoder_rep = torch.bmm(a, encoder_outputs)
  117.         weighted_encoder_rep = weighted_encoder_rep.permute(1, 0, 2)
  119.         return weighted_encoder_rep
  122.     def forward(self,
  123.                 input: Tensor,
  124.                 decoder_hidden: Tensor,
  125.                 encoder_outputs: Tensor) -> Tuple[Tensor]:
  127.         input = input.unsqueeze(0)
  129.         embedded = self.dropout(self.embedding(input))
  131.         weighted_encoder_rep = self._weighted_encoder_rep(decoder_hidden,
  132.                                                           encoder_outputs)
  134.         rnn_input = torch.cat((embedded, weighted_encoder_rep), dim = 2)
  136.         output, decoder_hidden = self.rnn(rnn_input, decoder_hidden.unsqueeze(0))
  138.         embedded = embedded.squeeze(0)
  139.         output = output.squeeze(0)
  140.         weighted_encoder_rep = weighted_encoder_rep.squeeze(0)
  142.         output = self.out(torch.cat((output,
  143.                                      weighted_encoder_rep,
  144.                                      embedded), dim = 1))
  146.         return output, decoder_hidden.squeeze(0)
  149. class Seq2Seq(nn.Module):
  150.     def __init__(self,
  151.                  encoder: nn.Module,
  152.                  decoder: nn.Module,
  153.                  device: torch.device):
  154.         super().__init__()
  156.         self.encoder = encoder
  157.         self.decoder = decoder
  158.         self.device = device
  160.     def forward(self,
  161.                 src: Tensor,
  162.                 trg: Tensor,
  163.                 teacher_forcing_ratio: float = 0.5) -> Tensor:
  165.         batch_size = src.shape[1]
  166.         max_len = trg.shape[0]
  167.         trg_vocab_size = self.decoder.output_dim
  169.         outputs = torch.zeros(max_len, batch_size, trg_vocab_size).to(self.device)
  171.         encoder_outputs, hidden = self.encoder(src)
  173.         # first input to the decoder is the <sos> token
  174.         output = trg[0,:]
  176.         for t in range(1, max_len):
  177.             output, hidden = self.decoder(output, hidden, encoder_outputs)
  178.             outputs[t] = output
  179.             teacher_force = random.random() < teacher_forcing_ratio
  180.             top1 = output.max(1)[1]
  181.             output = (trg[t] if teacher_force else top1)
  183.         return outputs
  186. INPUT_DIM = len(SRC.vocab)
  187. OUTPUT_DIM = len(TRG.vocab)
  188. # ENC_EMB_DIM = 256
  189. # DEC_EMB_DIM = 256
  190. # ENC_HID_DIM = 512
  191. # DEC_HID_DIM = 512
  192. # ATTN_DIM = 64
  193. # ENC_DROPOUT = 0.5
  194. # DEC_DROPOUT = 0.5
  196. ENC_EMB_DIM = 32
  197. DEC_EMB_DIM = 32
  198. ENC_HID_DIM = 64
  199. DEC_HID_DIM = 64
  200. ATTN_DIM = 8
  201. ENC_DROPOUT = 0.5
  202. DEC_DROPOUT = 0.5
  204. enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)
  206. attn = Attention(ENC_HID_DIM, DEC_HID_DIM, ATTN_DIM)
  208. dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)
  210. model = Seq2Seq(enc, dec, device).to(device)
  213. def init_weights(m: nn.Module):
  214.     for name, param in m.named_parameters():
  215.         if 'weight' in name:
  216.             nn.init.normal_(param.data, mean=0, std=0.01)
  217.         else:
  218.             nn.init.constant_(param.data, 0)
  221. model.apply(init_weights)
  223. optimizer = optim.Adam(model.parameters())
  226. def count_parameters(model: nn.Module):
  227.     return sum(p.numel() for p in model.parameters() if p.requires_grad)
  230. print(f'The model has {count_parameters(model):,} trainable parameters')

输出： The model has 1,856,685 trainable parameters

注：当计算模型分数尤其是翻译模型时，我们需要设置nn.CrossEntropyLoss 忽略padding。

  1. PAD_IDX = TRG.vocab.stoi['<pad>']
  3. criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)

最后，我们可以训练和评价模型：

  1. import math
  2. import time
  5. def train(model: nn.Module,
  6.           iterator: BucketIterator,
  7.           optimizer: optim.Optimizer,
  8.           criterion: nn.Module,
  9.           clip: float):
  11.     model.train()
  13.     epoch_loss = 0
  15.     for _, batch in enumerate(iterator):
  17.         src = batch.src
  18.         trg = batch.trg
  20.         optimizer.zero_grad()
  22.         output = model(src, trg)
  24.         output = output[1:].view(-1, output.shape[-1])
  25.         trg = trg[1:].view(-1)
  27.         loss = criterion(output, trg)
  29.         loss.backward()
  31.         torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
  33.         optimizer.step()
  35.         epoch_loss += loss.item()
  37.     return epoch_loss / len(iterator)
  40. def evaluate(model: nn.Module,
  41.              iterator: BucketIterator,
  42.              criterion: nn.Module):
  44.     model.eval()
  46.     epoch_loss = 0
  48.     with torch.no_grad():
  50.         for _, batch in enumerate(iterator):
  52.             src = batch.src
  53.             trg = batch.trg
  55.             output = model(src, trg, 0) #turn off teacher forcing
  57.             output = output[1:].view(-1, output.shape[-1])
  58.             trg = trg[1:].view(-1)
  60.             loss = criterion(output, trg)
  62.             epoch_loss += loss.item()
  64.     return epoch_loss / len(iterator)
  67. def epoch_time(start_time: int,
  68.                end_time: int):
  69.     elapsed_time = end_time - start_time
  70.     elapsed_mins = int(elapsed_time / 60)
  71.     elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
  72.     return elapsed_mins, elapsed_secs
  75. N_EPOCHS = 10
  76. CLIP = 1
  78. best_valid_loss = float('inf')
  80. for epoch in range(N_EPOCHS):
  82.     start_time = time.time()
  84.     train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
  85.     valid_loss = evaluate(model, valid_iterator, criterion)
  87.     end_time = time.time()
  89.     epoch_mins, epoch_secs = epoch_time(start_time, end_time)
  91.     print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
  92.     print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
  93.     print(f'\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')
  95. test_loss = evaluate(model, test_iterator, criterion)
  97. print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')

输出： Epoch: 01 | Time: 0m 36s
Train Loss: 5.686 | Train PPL: 294.579
Val. Loss: 5.250 | Val. PPL: 190.638
Epoch: 02 | Time: 0m 37s
Train Loss: 5.019 | Train PPL: 151.260
Val. Loss: 5.155 | Val. PPL: 173.274 Epoch: 03 | Time: 0m 37s
Train Loss: 4.757 | Train PPL: 116.453
Val. Loss: 4.976 | Val. PPL: 144.824
Epoch: 04 | Time: 0m 35s
Train Loss: 4.574 | Train PPL: 96.914
Val. Loss: 4.835 | Val. PPL: 125.834
Epoch: 05 | Time: 0m 35s
Train Loss: 4.421 | Train PPL: 83.185
Val. Loss: 4.783 | Val. PPL: 119.414
Epoch: 06 | Time: 0m 38s
Train Loss: 4.321 | Train PPL: 75.233
Val. Loss: 4.802 | Val. PPL: 121.734
Epoch: 07 | Time: 0m 38s
Train Loss: 4.233 | Train PPL: 68.957
Val. Loss: 4.675 | Val. PPL: 107.180
Epoch: 08 | Time: 0m 35s
Train Loss: 4.108 | Train PPL: 60.838
Val. Loss: 4.622 | Val. PPL: 101.693
Epoch: 09 | Time: 0m 34s
Train Loss: 4.020 | Train PPL: 55.680
Val. Loss: 4.530 | Val. PPL: 92.785 Epoch: 10 | Time: 0m 34s
Train Loss: 3.919 | Train PPL: 50.367
Val. Loss: 4.448 | Val. PPL: 85.441
| Test Loss: 4.464 | Test PPL: 86.801 |

接下来的步骤

  • 看看Ben Trevett使用torchtext教程的其余部分
  • 请继续关注使用其他torchtext功能以及nn.Transformer语言建模预测下一个单词的教程！

脚本的总运行时间： (6分钟27.732秒）