3 - Neural Machine Translation by Jointly Learning to Align and Translate
В этом третьем посте о моделях sequence-to-sequence с использованием PyTorch и torchText мы будем реализовывать модель из стать Neural Machine Translation by Jointly Learning to Align and Translate. Эта модель демонстрирует лучшую точность из из трёх моделей (~27 по сравнению с ~34 у предыдущей модели).
Как и ранее, если визуальный формат поста вас не удовлетворяет, то ниже ссылки на английскую и русскую версию jupyter notebook:
Исходная версия (Open jupyter notebook In Colab)
Русская версия (Open jupyter notebook In Colab)
Введение
Напоминаем общую модель кодера-декодера:
![](https://habrastorage.org/getpro/habr/upload_files/2ea/527/37f/2ea52737f8b23db2c21775bfefb6bbdc.png)
В предыдущей модели наша архитектура была построена таким образом, чтобы уменьшить «сжатие информации» путем явной передачи вектора контекста в декодер на линейный слой
на каждом временном шаге, совместно с передачей входного слова, прошедшего через слой эмбеддинга,
и со скрытым состоянием
.
![](https://habrastorage.org/getpro/habr/upload_files/8cd/1cb/6fa/8cd1cb6fa9090497ea69f677c389aa13.png)
Несмотря на то, что мы частично уменьшили сжатие информации, наш вектор контекста по-прежнему должен содержать всю информацию об исходном предложении. Модель, реализованная в этом разделе, избегает такого сжатия, позволяя декодеру просматривать все исходное предложение черезегоскрытыесостояния на каждом этапе декодирования! Как это стало возможным? Благодаря вниманию.
Для использования механизма внимания, сначала вычисляем вектор внимания . Каждый элемент вектора внимания находится в диапазоне от 0 до 1, а сумма элементов вектора равна 1. Затем мы вычисляем взвешенную сумму скрытых состояний исходного предложения
, чтобы получить взвешенный исходный вектор
.
Мы вычисляем новый взвешенный исходный вектор на каждом временном шаге при декодировании, используя его в качестве входных данных для RNN декодера, а также линейного слоя для прогнозирования. Мы объясним, как все это сделать далее.
Подготовка данных
Снова подготовка аналогична прошлой.
Сначала мы импортируем все необходимые модули.
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchtext.legacy.datasets import Multi30k
from torchtext.legacy.data import Field, BucketIterator
import spacy
import numpy as np
import random
import math
import time
Установите случайные значения для воспроизводимости.
SEED = 1234
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
Загрузите немецкую и английскую модели spaCy.
python -m spacy download en_core_web_sm
python -m spacy download de_core_news_sm
Для загрузки в Google Colab используем следующие команды (После загрузки обязательно перезапустите colab runtime! Наибыстрейший способ через короткую комаду: Ctrl + M + .):
!pip install -U spacy==3.0
!python -m spacy download en_core_web_sm
!python -m spacy download de_core_news_sm
spacy_de = spacy.load('de_core_news_sm')
spacy_en = spacy.load('en_core_web_sm')
Создаем токенизаторы.
def tokenize_de(text):
"""
Tokenizes German text from a string into a list of strings
"""
return [tok.text for tok in spacy_de.tokenizer(text)]
def tokenize_en(text):
"""
Tokenizes English text from a string into a list of strings
"""
return [tok.text for tok in spacy_en.tokenizer(text)]
Поля остаются теми же, что и раньше.
SRC = Field(tokenize = tokenize_de,
init_token = '<sos>',
eos_token = '<eos>',
lower = True)
TRG = Field(tokenize = tokenize_en,
init_token = '<sos>',
eos_token = '<eos>',
lower = True)
Загружаем данные.
train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'),
fields = (SRC, TRG))
Создаём словари.
SRC.build_vocab(train_data, min_freq = 2)
TRG.build_vocab(train_data, min_freq = 2)
Определяем устройство.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Создаём итераторы.
BATCH_SIZE = 128
train_iterator, valid_iterator, test_iterator = BucketIterator.splits(
(train_data, valid_data, test_data),
batch_size = BATCH_SIZE,
device = device)
Создание модели Seq2Seq
Кодер
Сначала мы создадим кодер. Как и в предыдущей модели, мы используем только один слой GRU, однако теперь он будет иметь вид двунаправленной RNN. В двунаправленной RNN у нас есть две RNN на каждом уровне. Вперёд-направленная RNN перебирает предложение, прошедшее эмбеддинга, слева направо показанонижезеленымцветом, а назад-направленная RNN перебирает предложение, прошедшее эмбеддинг, справа налево бирюзовый. Все, что нам нужно сделать в коде — установить bidirectional = True
, а затем провести предложение через слой эмбеддинга в RNN, как и раньше.
![](https://habrastorage.org/getpro/habr/upload_files/b24/32e/e39/b2432ee39f34e971a66baacb43ac6586.png)
Теперь у нас есть:
Где и
.
Как и раньше, мы передаем в RNN только ввод (embedded
), который сообщает PyTorch о необходимости инициализировать как прямое, так и обратное начальные скрытые состояния ( and
, respectively) тензором с нулевыми значениями элементов. Кроме того, мы получаем два вектора контекста: один из прямой RNN после того, как она увидит последнее слово в предложении
, а второй из обратной RNN после того, как она зафиксирует первое слово в предложении
.
RNN возвращает outputs
и hidden
.
outputs
имеет размер srclen,batchsize,hiddim∗numdirections где первые hid_dim
элементов в третьем измерении - это скрытые состояния от верхнего уровня вперёд-направленной RNN, а последнее hid_dim
элементов — это скрытые состояния от верхнего уровня назад-направленной RNN. Мы можем думать о третьем измерении как о прямом и обратном скрытых состояниях, связанных вместе друг с другом, т.е. ,
и мы можем обозначить все скрытые состояния кодировщика прямоеиобратноесцеплениевместе как
тензорконтекста.
hidden
имеет размер nlayers∗numdirections,batchsize,hiddim, где −2,:,: дает скрытое состояние вперёд-направленной RNN верхнего уровня после последнего временного шага т.е.послетого,каконувиделпоследнееслововпредложении и −1,:,: дает верхнему уровню скрытое состояние обратно-направленной RNN после последнего временного шага т.е.послетого,каконувиделпервоеслововпредложении.
Поскольку декодер не является двунаправленным, ему нужен только один вектор контекста для использования в качестве начального скрытого состояния
, но в настоящее время у нас есть два вектора контекста (
и
, respectively). Мы решаем эту проблему, объединив два вектора контекста вместе, пропустив их через линейный слой
и применяя функцию активации
.
Замечание: на самом деле здесь есть некоторое отклонение от реализации в статье. В статье авторы передают только первое назад-направленной скрытое состояние RNN через линейный слой, чтобы получить начальное скрытое состояние вектора контекста для декодера. Это кажется бессмысленным, поэтому мы изменили эту часть формирования вектора внимания.
Поскольку мы хотим, чтобы наша модель просматривала все исходное предложение, мы возвращаем outputs
, в виде объединённых скрытых состояний вперед и назад для каждого токена в исходном предложении. Мы возвращаем hidden
, который действует как начальное скрытое состояние в декодере.
class Encoder(nn.Module):
def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
super().__init__()
self.embedding = nn.Embedding(input_dim, emb_dim)
self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, src):
#src = [src len, batch size]
embedded = self.dropout(self.embedding(src))
#embedded = [src len, batch size, emb dim]
outputs, hidden = self.rnn(embedded)
#outputs = [src len, batch size, hid dim * num directions]
#hidden = [n layers * num directions, batch size, hid dim]
#hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]
#outputs are always from the last layer
#hidden [-2, :, : ] is the last of the forwards RNN
#hidden [-1, :, : ] is the last of the backwards RNN
#initial decoder hidden is final hidden state of the forwards and backwards
# encoder RNNs fed through a linear layer
hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
#outputs = [src len, batch size, enc hid dim * 2]
#hidden = [batch size, dec hid dim]
return outputs, hidden
Внимание
Далее идет слой внимания. Этот слой принимает предыдущее скрытое состояние декодера и все скрытые состояния кодера, собранные в тензор контекста
. Слой генерирует вектор внимания
длины исходного предложения, каждый элемент которого находится в диапазоне от 0 до 1, а вся сумма элементов вектора равна 1.
Интуитивно понятно, что этот слой берет то, что мы уже декодировали, , и все, что мы закодировали в
, для создания вектора
, который представляет, каким словам в исходном предложении мы должны уделять большее внимание для правильного предсказать следующее слова декодировщиком,
.
Сначала мы вычисляем энергию взаимодействия между предыдущим скрытым состоянием декодера и скрытыми состояниями кодера. Поскольку скрытые состояния нашего кодера представляют собой последовательность тензоров, и наше предыдущее скрытое состояние декодера — это одиночный тензор, первое, что мы делаем, это
повторяем
предыдущее скрытое состояние декодера раз. Затем мы вычисляем энергию взаимодействия
между ними, объединив их вместе и пропустив через линейный слой (
attn
) и функцию активации .
Эту величину можно рассматривать как вычисление того, насколько хорошо каждое скрытое состояние кодера «совпадает» с предыдущим скрытым состоянием декодера.
В настоящее время у нас есть dec hid dim, src len тензор для каждого примера в батче. Мы хотим, чтобы он был длины src len для каждого примера в батче, так как внимание должно быть длины исходного предложения. Это достигается путем умножения энергии
на 1, dec hid dim-размерный тезор .
Мы можем думать о как о качестве весов взвешенной суммы энергии по всем скрытым состояниям кодировщика. Эти веса говорят нам, насколько мы должны уделять внимание каждому токену в исходной последовательности. Параметры
инициализируются случайным образом, но изучаются вместе с остальной частью модели посредством обратного распространения ошибки. Обратите внимание, как v не зависит от времени, и то же время
используется для каждого временного шага декодирования. Реализуем
как линейный слой без смещения.
Наконец, мы следим за тем, чтобы вектор внимания соответствовал ограничениям, накладываемым на элементы этого вектора при передаче его через слой : все элементы находятся между 0 и 1, и суммирование элементов даёт 1.
Это привлекает внимание к исходному предложению!
Графически это выглядит примерно так, как показано ниже. Так для вычисления самого первого вектора внимания . Зеленые блоки представляют скрытые состояния как от вперёд-направленной, так и назад-направленной RNN, и все вычисления внимания выполняются в розовом блоке.
![](https://habrastorage.org/getpro/habr/upload_files/bf1/c96/a01/bf1c96a01f192c7b03516afd32f95998.png)
class Attention(nn.Module):
def __init__(self, enc_hid_dim, dec_hid_dim):
super().__init__()
self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
self.v = nn.Linear(dec_hid_dim, 1, bias = False)
def forward(self, hidden, encoder_outputs):
#hidden = [batch size, dec hid dim]
#encoder_outputs = [src len, batch size, enc hid dim * 2]
batch_size = encoder_outputs.shape[1]
src_len = encoder_outputs.shape[0]
#repeat decoder hidden state src_len times
hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
encoder_outputs = encoder_outputs.permute(1, 0, 2)
#hidden = [batch size, src len, dec hid dim]
#encoder_outputs = [batch size, src len, enc hid dim * 2]
energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim = 2)))
#energy = [batch size, src len, dec hid dim]
attention = self.v(energy).squeeze(2)
#attention= [batch size, src len]
return F.softmax(attention, dim=1)
Декодер
Далее идет декодер.
Декодер содержит слой внимания, attention
, который принимает предыдущее скрытое состояние , все скрытые состояния кодировщика
, и возвращает вектор внимания
.
Затем мы используем этот вектор внимания для создания взвешенного исходного вектора , который обозначается как
weighted
, который представляет собой взвешенную сумму скрытых состояний кодировщика , использованный совместно с весами
.
Входное слово, прошедшее эмбеддинга , взвешенный исходный вектор
, и предыдущее скрытое состояние декодера
, все это передаются в декодер RNN, с
и
и соединяется вместе.
Затем мы передаем ,
и
через линейный слой
для совершения предсказания следующего слова в целевом предложении
. Это делается путем их объединения.
На изображении ниже показано декодирование первого слова в примере перевода.
![](https://habrastorage.org/getpro/habr/upload_files/e34/da6/b31/e34da6b31765b7c156b81de3006cc723.png)
Зелёные/бирюзовый блоки показывают RNNs кодера которые выдают H, красный блок показывает вектор контекста, , синий блок показывает RNN декодера, который выводит s_t, фиолетовый блок показывает линейный слой
, выводит
, а оранжевый блок показывает вычисление взвешенной суммы по
от
и выходов
. Не показан расчет
.
class Decoder(nn.Module):
def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
super().__init__()
self.output_dim = output_dim
self.attention = attention
self.embedding = nn.Embedding(output_dim, emb_dim)
self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, input, hidden, encoder_outputs):
#input = [batch size]
#hidden = [batch size, dec hid dim]
#encoder_outputs = [src len, batch size, enc hid dim * 2]
input = input.unsqueeze(0)
#input = [1, batch size]
embedded = self.dropout(self.embedding(input))
#embedded = [1, batch size, emb dim]
a = self.attention(hidden, encoder_outputs)
#a = [batch size, src len]
a = a.unsqueeze(1)
#a = [batch size, 1, src len]
encoder_outputs = encoder_outputs.permute(1, 0, 2)
#encoder_outputs = [batch size, src len, enc hid dim * 2]
weighted = torch.bmm(a, encoder_outputs)
#weighted = [batch size, 1, enc hid dim * 2]
weighted = weighted.permute(1, 0, 2)
#weighted = [1, batch size, enc hid dim * 2]
rnn_input = torch.cat((embedded, weighted), dim = 2)
#rnn_input = [1, batch size, (enc hid dim * 2) + emb dim]
output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
#output = [seq len, batch size, dec hid dim * n directions]
#hidden = [n layers * n directions, batch size, dec hid dim]
#seq len, n layers and n directions will always be 1 in this decoder, therefore:
#output = [1, batch size, dec hid dim]
#hidden = [1, batch size, dec hid dim]
#this also means that output == hidden
assert (output == hidden).all()
embedded = embedded.squeeze(0)
output = output.squeeze(0)
weighted = weighted.squeeze(0)
prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))
#prediction = [batch size, output dim]
return prediction, hidden.squeeze(0)
Seq2Seq
Это первая модель, в которой нам не нужно, чтобы RNN кодировщика и RNN декодера имели одинаковые скрытые размеры, однако кодировщик должен быть двунаправленным. Последнее требование можно игнорировать, изменив все размерность входных данных с enc_dim * 2
на enc_dim * 2 if encoder_is_bidirectional else enc_dim
.
Эта модель seq2seq инкапсулирует кодер и декодер как и в двух предыдущих моделях. Единственная разница в том, что encoder
возвращает как окончательное скрытое состояние который является окончательным скрытым состоянием как от вперёд-направленного, так и от назад-направленного RNN кодировщика, прошедших через линейный уровень для использования в качестве начального скрытого состояния в декодере, а также для каждого скрытого состояния которые представляют собой скрытые состояния на выходе вперёд- и назад-направленные RNNN, накладываемые друг на друга. Нам также необходимо обеспечить, чтобы hidden
и encoder_outputs
передавались в декодер.
Кратко пройдемся по всем этапам:
тензор
outputs
создан для хранения всех прогнозовисходная последовательность
, подается в кодировщик для получения
и
начальное скрытое состояние декодера установлено как вектор
context
мы используем батч токенов
<sos>
как первыйinput
затем декодируем в цикле:
вставка входного токена
, предыдущее скрытое состояние
, и все выходы кодера
в декодер
получение прогноза
и новое скрытое состояние
затем мы решаем, собираемся ли мы применять обучение с принуждением или нет, устанавливая следующий ввод соответствующим образом
class Seq2Seq(nn.Module):
def __init__(self, encoder, decoder, device):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.device = device
def forward(self, src, trg, teacher_forcing_ratio = 0.5):
#src = [src len, batch size]
#trg = [trg len, batch size]
#teacher_forcing_ratio is probability to use teacher forcing
#e.g. if teacher_forcing_ratio is 0.75 we use teacher forcing 75% of the time
batch_size = src.shape[1]
trg_len = trg.shape[0]
trg_vocab_size = self.decoder.output_dim
#tensor to store decoder outputs
outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)
#encoder_outputs is all hidden states of the input sequence, back and forwards
#hidden is the final forward and backward hidden states, passed through a linear layer
encoder_outputs, hidden = self.encoder(src)
#first input to the decoder is the <sos> tokens
input = trg[0,:]
for t in range(1, trg_len):
#insert input token embedding, previous hidden state and all encoder hidden states
#receive output tensor (predictions) and new hidden state
output, hidden = self.decoder(input, hidden, encoder_outputs)
#place predictions in a tensor holding predictions for each token
outputs[t] = output
#decide if we are going to use teacher forcing or not
teacher_force = random.random() < teacher_forcing_ratio
#get the highest predicted token from our predictions
top1 = output.argmax(1)
#if teacher forcing, use actual next token as next input
#if not, use predicted token
input = trg[t] if teacher_force else top1
return outputs
Обучение модели Seq2Seq
Остальная часть этого урока очень похожа на предыдущий.
Мы инициализируем наши параметры, кодера, декодер и модели seq2seq поместив его на графический процессор, если он у нас есть.
INPUT_DIM = len(SRC.vocab)
OUTPUT_DIM = len(TRG.vocab)
ENC_EMB_DIM = 256
DEC_EMB_DIM = 256
ENC_HID_DIM = 512
DEC_HID_DIM = 512
ENC_DROPOUT = 0.5
DEC_DROPOUT = 0.5
attn = Attention(ENC_HID_DIM, DEC_HID_DIM)
enc = Encoder(INPUT_DIM, ENC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, ENC_DROPOUT)
dec = Decoder(OUTPUT_DIM, DEC_EMB_DIM, ENC_HID_DIM, DEC_HID_DIM, DEC_DROPOUT, attn)
model = Seq2Seq(enc, dec, device).to(device)
Мы используем упрощенную версию схемы инициализации весов, использованную в статье. Здесь мы инициализируем все смещения равными нулю и все веса из .
def init_weights(m):
for name, param in m.named_parameters():
if 'weight' in name:
nn.init.normal_(param.data, mean=0, std=0.01)
else:
nn.init.constant_(param.data, 0)
model.apply(init_weights)
Подсчитаем количество параметров. Получаем прибавку почти 50% по сравнению с количеством параметров из последней модели.
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'The model has {count_parameters(model):,} trainable parameters')
Создаем оптимизатор.
optimizer = optim.Adam(model.parameters())
Инициализируем функцию потерь.
TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]
criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)
Затем мы создаем цикл обучения ...
def train(model, iterator, optimizer, criterion, clip):
model.train()
epoch_loss = 0
for i, batch in enumerate(iterator):
src = batch.src
trg = batch.trg
optimizer.zero_grad()
output = model(src, trg)
#trg = [trg len, batch size]
#output = [trg len, batch size, output dim]
output_dim = output.shape[-1]
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
#trg = [(trg len - 1) * batch size]
#output = [(trg len - 1) * batch size, output dim]
loss = criterion(output, trg)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
optimizer.step()
epoch_loss += loss.item()
return epoch_loss / len(iterator)
...и цикл оценки, не забывая установить модель на eval
режим и отключив обучение с принуждением.
def evaluate(model, iterator, criterion):
model.eval()
epoch_loss = 0
with torch.no_grad():
for i, batch in enumerate(iterator):
src = batch.src
trg = batch.trg
output = model(src, trg, 0) #turn off teacher forcing
#trg = [trg len, batch size]
#output = [trg len, batch size, output dim]
output_dim = output.shape[-1]
output = output[1:].view(-1, output_dim)
trg = trg[1:].view(-1)
#trg = [(trg len - 1) * batch size]
#output = [(trg len - 1) * batch size, output dim]
loss = criterion(output, trg)
epoch_loss += loss.item()
return epoch_loss / len(iterator)
Наконец, определим функцию подсчёта времени.
def epoch_time(start_time, end_time):
elapsed_time = end_time - start_time
elapsed_mins = int(elapsed_time / 60)
elapsed_secs = int(elapsed_time - (elapsed_mins * 60))
return elapsed_mins, elapsed_secs
Затем мы обучаем нашу модель, сохраняя параметры, которые дают нам наименьшие потери при проверке.
N_EPOCHS = 10
CLIP = 1
best_valid_loss = float('inf')
for epoch in range(N_EPOCHS):
start_time = time.time()
train_loss = train(model, train_iterator, optimizer, criterion, CLIP)
valid_loss = evaluate(model, valid_iterator, criterion)
end_time = time.time()
epoch_mins, epoch_secs = epoch_time(start_time, end_time)
if valid_loss < best_valid_loss:
best_valid_loss = valid_loss
torch.save(model.state_dict(), 'tut3-model.pt')
print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')
print(f'\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')
print(f'\t Val. Loss: {valid_loss:.3f} | Val. PPL: {math.exp(valid_loss):7.3f}')
Наконец, мы тестируем модель на тестовой выборке, используя эти «лучшие» параметры.
model.load_state_dict(torch.load('tut3-model.pt'))
test_loss = evaluate(model, test_iterator, criterion)
print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')
Мы улучшили предыдущую модель, но это произошло за счет удвоения времени обучения.
В следующем разделе мы будем использовать ту же архитектуру, но применим несколько приемов ко всем архитектурам RNN - упакованные дополненные последовательности и маскирование. Мы реализуем код, который позволит нам посмотреть, на какие слова во входных данных RNN обращает внимание при декодировании выходных данных.