15.10. Предобучение BERT¶
Имея реализованную модель BERT Section 15.8 и примеры для предобучения, полученные из набора данных WikiText-2 Section 15.9, в этом разделе мы выполним предобучение BERT на наборе данных WikiText-2.
import torch
from torch import nn
from d2l import torch as d2l
Для начала загрузим набор данных WikiText-2 в виде мини-батчей примеров для предобучения, предназначенных для маскированного языкового моделирования и предсказания следующего предложения. Размер батча равен 512, а максимальная длина входной последовательности BERT составляет 64. Обратите внимание, что в оригинальной модели BERT максимальная длина равна 512.
batch_size, max_len = 512, 64
train_iter, vocab = d2l.load_data_wiki(batch_size, max_len)
15.10.1. Предобучение BERT¶
Оригинальная модель BERT имеет две версии разного размера (Devlin et al., 2018). Базовая модель (\(\textrm{BERT}_{\textrm{BASE}}\)) использует 12 слоёв (блоков кодировщика трансформера) с 768 скрытыми единицами (скрытый размер) и 12 головок самовнимания. Большая модель (\(\textrm{BERT}_{\textrm{LARGE}}\)) использует 24 слоя с 1024 скрытыми единицами и 16 головок самовнимания. Примечательно, что первая содержит 110 миллионов параметров, а вторая — 340 миллионов. Для удобства демонстрации мы определим маленький BERT, использующий 2 слоя, 128 скрытых единиц и 2 головки самовнимания.
net = d2l.BERTModel(len(vocab), num_hiddens=128,
ffn_num_hiddens=256, num_heads=2, num_blks=2, dropout=0.2)
devices = d2l.try_all_gpus()
loss = nn.CrossEntropyLoss()
Прежде чем определить цикл обучения, мы введем вспомогательную функцию
_get_batch_loss_bert. Принимая на вход фрагмент обучающих примеров,
эта функция вычисляет функцию потерь для обеих задач: маскированного
языкового моделирования и предсказания следующего предложения. Обратите
внимание, что итоговая функция потерь при предобучении BERT представляет
собой сумму функций потерь маскированного языкового моделирования и
предсказания следующего предложения.
#@save
def _get_batch_loss_bert(net, loss, vocab_size, tokens_X,
segments_X, valid_lens_x,
pred_positions_X, mlm_weights_X,
mlm_Y, nsp_y):
# Forward pass
_, mlm_Y_hat, nsp_Y_hat = net(tokens_X, segments_X,
valid_lens_x.reshape(-1),
pred_positions_X)
# Compute masked language model loss
mlm_l = loss(mlm_Y_hat.reshape(-1, vocab_size), mlm_Y.reshape(-1)) *\
mlm_weights_X.reshape(-1, 1)
mlm_l = mlm_l.sum() / (mlm_weights_X.sum() + 1e-8)
# Compute next sentence prediction loss
nsp_l = loss(nsp_Y_hat, nsp_y)
l = mlm_l + nsp_l
return mlm_l, nsp_l, l
Используя две вышеупомянутые вспомогательные функции, следующая
train_bert функция определяет процедуру предобучения BERT (net)
на наборе данных WikiText-2 (train_iter). Предобучение BERT может
занимать очень много времени. Вместо указания количества эпох для
обучения, как в функции train_ch13 (см.
Section 14.1), входной параметр num_steps
следующей функции указывает количество итераций для обучения.
def train_bert(train_iter, net, loss, vocab_size, devices, num_steps):
net(*next(iter(train_iter))[:4])
net = nn.DataParallel(net, device_ids=devices).to(devices[0])
trainer = torch.optim.Adam(net.parameters(), lr=0.01)
step, timer = 0, d2l.Timer()
animator = d2l.Animator(xlabel='step', ylabel='loss',
xlim=[1, num_steps], legend=['mlm', 'nsp'])
# Sum of masked language modeling losses, sum of next sentence prediction
# losses, no. of sentence pairs, count
metric = d2l.Accumulator(4)
num_steps_reached = False
while step < num_steps and not num_steps_reached:
for tokens_X, segments_X, valid_lens_x, pred_positions_X,\
mlm_weights_X, mlm_Y, nsp_y in train_iter:
tokens_X = tokens_X.to(devices[0])
segments_X = segments_X.to(devices[0])
valid_lens_x = valid_lens_x.to(devices[0])
pred_positions_X = pred_positions_X.to(devices[0])
mlm_weights_X = mlm_weights_X.to(devices[0])
mlm_Y, nsp_y = mlm_Y.to(devices[0]), nsp_y.to(devices[0])
trainer.zero_grad()
timer.start()
mlm_l, nsp_l, l = _get_batch_loss_bert(
net, loss, vocab_size, tokens_X, segments_X, valid_lens_x,
pred_positions_X, mlm_weights_X, mlm_Y, nsp_y)
l.backward()
trainer.step()
metric.add(mlm_l, nsp_l, tokens_X.shape[0], 1)
timer.stop()
animator.add(step + 1,
(metric[0] / metric[3], metric[1] / metric[3]))
step += 1
if step == num_steps:
num_steps_reached = True
break
print(f'MLM loss {metric[0] / metric[3]:.3f}, '
f'NSP loss {metric[1] / metric[3]:.3f}')
print(f'{metric[2] / timer.sum():.1f} sentence pairs/sec on '
f'{str(devices)}')
Мы можем построить графики функции потерь маскированного языкового моделирования и функции потерь предсказания следующего предложения в процессе предобучения BERT.
train_bert(train_iter, net, loss, len(vocab), devices, 50)
15.10.2. Представление текста с помощью BERT¶
После предобучения BERT мы можем использовать его для представления
отдельных текстов, пар текстов или любого токена в них. Следующая
функция возвращает представления BERT (net) для всех токенов в
tokens_a и tokens_b.
def get_bert_encoding(net, tokens_a, tokens_b=None):
tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b)
token_ids = torch.tensor(vocab[tokens], device=devices[0]).unsqueeze(0)
segments = torch.tensor(segments, device=devices[0]).unsqueeze(0)
valid_len = torch.tensor(len(tokens), device=devices[0]).unsqueeze(0)
encoded_X, _, _ = net(token_ids, segments, valid_len)
return encoded_X
Рассмотрим предложение «a crane is flying». Вспомним представление
входных данных BERT, обсуждавшееся в Section 15.8.4.
После вставки специальных токенов «<cls>» (используемого для
классификации) и «<sep>» (используемого для разделения) длина входной
последовательности BERT составляет шесть. Поскольку нуль — это индекс
токена «<cls>», encoded_text[:, 0, :] представляет собой
представление всего входного предложения в BERT. Чтобы проиллюстрировать
работу с многозначным токеном «crane», мы также выводим первые три
элемента представления этого токена в BERT.
tokens_a = ['a', 'crane', 'is', 'flying']
encoded_text = get_bert_encoding(net, tokens_a)
# Tokens: '<cls>', 'a', 'crane', 'is', 'flying', '<sep>'
encoded_text_cls = encoded_text[:, 0, :]
encoded_text_crane = encoded_text[:, 2, :]
encoded_text.shape, encoded_text_cls.shape, encoded_text_crane[0][:3]
Теперь рассмотрим пару предложений: «a crane driver came» и «he just
left». Аналогично, encoded_pair[:, 0, :] представляет собой
закодированный результат всей пары предложений, полученный с помощью
предварительно обученной модели BERT. Обратите внимание, что первые три
элемента многозначного токена «crane» отличаются от тех, что были
получены при другом контексте. Это подтверждает, что представления BERT
чувствительны к контексту.
tokens_a, tokens_b = ['a', 'crane', 'driver', 'came'], ['he', 'just', 'left']
encoded_pair = get_bert_encoding(net, tokens_a, tokens_b)
# Tokens: '<cls>', 'a', 'crane', 'driver', 'came', '<sep>', 'he', 'just',
# 'left', '<sep>'
encoded_pair_cls = encoded_pair[:, 0, :]
encoded_pair_crane = encoded_pair[:, 2, :]
encoded_pair.shape, encoded_pair_cls.shape, encoded_pair_crane[0][:3]
В Section 16 мы будем дообучать предварительно обученную модель BERT для задач обработки естественного языка.
15.10.3. Итоги¶
Оригинальная модель BERT представлена в двух версиях: базовая модель содержит 110 миллионов параметров, а большая — 340 миллионов параметров.
После предобучения BERT мы можем использовать его для представления отдельных текстов, пар текстов или любых токенов в них.
В ходе эксперимента один и тот же токен имеет различные представления BERT в зависимости от контекста. Это подтверждает, что представления BERT чувствительны к контексту.
15.10.4. Упражнения¶
В ходе эксперимента мы видим, что функция потерь маскированного языкового моделирования значительно выше, чем функция потерь предсказания следующего предложения. Почему?
Установите максимальную длину входной последовательности BERT равной 512 (как в оригинальной модели BERT). Используйте конфигурации оригинальной модели BERT, такие как \(\textrm{BERT}_{\textrm{LARGE}}\). Возникает ли какая-либо ошибка при запуске этого раздела? Почему?