Внимание — это все, что вам нужно (II) — изучение кода
Поскольку теоретическая часть этой статьи была рассмотрена в этом посте, в этом посте я сосредоточусь на реализации каждого компонента, упомянутого в статье Transformer.
Первоначальные авторы уже опубликовали свой код реализации здесь. Тем не менее, эти почти 3000 строк кода на Python не так удобны для новичка (или, по крайней мере, не так много для человека с опытом работы в CS), как я.
Keras был моей основной средой программирования ML с первого дня. Естественно, я бы попытался использовать Keras для реализации Transformer. Но очевидно, что в Keras еще нет предопределенных слоев для мультиголовки, самоконтроля или позиционного кодирования. Это в значительной степени означает, что мне придется реализовать эти слои самостоятельно, используя определяемые пользователем слои Keras.
После нескольких часов кодирования я вскоре понял, что я определенно не первый, кто пытается кодировать Transformer с помощью Keras, и, возможно, уже есть некоторые реализации, доступные в Интернете. К счастью, я нашел пару реализаций Transformer от Keras на PaperWithCode. PaperWithCode — хороший справочный веб-сайт, на котором собрано множество исследовательских работ и соответствующих реализаций.
Я беру эту версию реализации Keras для изучения кода, так как это, вероятно, самая лаконичная и чистая реализация Keras, которую я когда-либо встречал. Я взял этот фрагмент кода и запустил его в Google Colab со стандартной задачей перевода с английского на немецкий, аналогично тому, что упоминается в документе.
Давайте рассмотрим архитектуру Transformer, упомянутую в статье:
Как мы видим на этом рисунке, у нас есть следующие компоненты внутри трансформатора:
- Слой встраивания ввода/вывода
- Уровень позиционного кодирования
- Слой внимания с несколькими головками
- Добавление и нормализация слоев
- Слой прямой связи
В следующих абзацах я расскажу, как реализовать эти 5 компонентов с помощью фреймворка Keras. В этом посте я сосредоточусь на изучении каждого компонента, а не на сборе всего вместе. Полную реализацию вместе с процессами обучения/тестирования можно найти на оригинальном авторском Github. Я также преобразовал исходный код в блокнот ipynb в моем Github для изучения.
Прежде чем мы начнем изучать коды реализации, давайте посмотрим, какие фреймворки и библиотеки используются:
import os, sys, time, random import h5py import numpy as np from keras.models import * from keras.layers import * from keras.callbacks import * from keras.initializers import * import tensorflow as tf
Слой встраивания ввода/вывода
Для этого компонента используется слой Keras Embedding по умолчанию. В документации Keras есть подробное объяснение и исходный код для этого компонента по умолчанию.
Вот краткий код:
class Transformer:
def __init__(self, ..., d_model=256, ...):
...
d_emb = d_model
self.i_word_emb = Embedding(i_tokens.num(), d_emb)
...
def compile(self, optimizer='adam', active_layers=999):
src_seq_input = Input(shape=(None,), dtype='int32')
tgt_seq_input = Input(shape=(None,), dtype='int32')
src_seq = src_seq_input
tgt_seq = Lambda(lambda x:x[:,:-1])(tgt_seq_input)
tgt_true = Lambda(lambda x:x[:,1:])(tgt_seq_input)
src_emb = self.i_word_emb(src_seq)
tgt_emb = self.o_word_emb(tgt_seq)
if self.pos_emb:
src_emb = add_layer([src_emb, self.pos_emb(src_seq)])
tgt_emb = add_layer([tgt_emb, self.pos_emb(tgt_seq)])
src_emb = self.emb_dropout(src_emb)
...
Уровень позиционного кодирования
В документе слои позиционного кодирования добавляются непосредственно к выходным данным входных/выходных слоев внедрения. Следовательно, в коде мы можем видеть прямую операцию добавления, определенную в классе Transformer следующим образом:
class Transformer:
def __init__(self, ...):
...
self.pos_emb = PosEncodingLayer(len_limit, d_emb) if self.src_loc_info else None
self.emb_dropout = Dropout(dropout)
self.i_word_emb = Embedding(i_tokens.num(), d_emb)
if share_word_emb:
assert i_tokens.num() == o_tokens.num()
self.o_word_emb = i_word_emb
else: self.o_word_emb = Embedding(o_tokens.num(), d_emb)
...
def compile(self, optimizer='adam', active_layers=999):
...
src_emb = self.i_word_emb(src_seq)
tgt_emb = self.o_word_emb(tgt_seq)
if self.pos_emb:
src_emb = add_layer([src_emb, self.pos_emb(src_seq)])
tgt_emb = add_layer([tgt_emb, self.pos_emb(tgt_seq)])
А класс PosEncodingLayer() определяется следующим образом:
class PosEncodingLayer:
def __init__(self, max_len, d_emb):
self.pos_emb_matrix = Embedding(max_len, d_emb, trainable=False, weights=[GetPosEncodingMatrix(max_len, d_emb)])
def get_pos_seq(self, x):
mask = K.cast(K.not_equal(x, 0), 'int32')
pos = K.cumsum(K.ones_like(x, 'int32'), 1)
return pos * mask
def __call__(self, seq, pos_input=False):
x = seq
if not pos_input: x = Lambda(self.get_pos_seq)(x)
return self.pos_emb_matrix(x)
Слой внимания с несколькими головками
Это, пожалуй, самый важный слой в модели Transformer. В документе архитектура выглядит так, как показано на рисунке ниже:
Как показано на рисунке, было бы естественно подумать о создании масштабируемого модуля внимания с точечным произведением и складывать его несколько раз, чтобы добиться внимания с несколькими головками. Эту идею можно найти, как показано ниже:
class ScaledDotProductAttention():
def __init__(self, attn_dropout=0.1):
self.dropout = Dropout(attn_dropout)
def __call__(self, q, k, v, mask): # mask_k or mask_qk
temper = tf.sqrt(tf.cast(tf.shape(k)[-1], dtype='float32'))
attn = Lambda(lambda x:K.batch_dot(x[0],x[1],axes=[2,2])/temper)([q, k]) # shape=(batch, q, k)
if mask is not None:
mmask = Lambda(lambda x:(-1e+9)*(1.-K.cast(x, 'float32')))(mask)
attn = Add()([attn, mmask])
attn = Activation('softmax')(attn)
attn = self.dropout(attn)
output = Lambda(lambda x:K.batch_dot(x[0], x[1]))([attn, v])
return output, attn
Теперь, когда этот класс ScaledDotProductAttention() определен, мы можем сложить этот модуль, чтобы сформировать внимание с несколькими головками:
class MultiHeadAttention():
# mode 0 - big martixes, faster; mode 1 - more clear implementation
def __init__(self, n_head, d_model, dropout, mode=0):
self.mode = mode
self.n_head = n_head
self.d_k = self.d_v = d_k = d_v = d_model // n_head
self.dropout = dropout
if mode == 0:
self.qs_layer = Dense(n_head*d_k, use_bias=False)
self.ks_layer = Dense(n_head*d_k, use_bias=False)
self.vs_layer = Dense(n_head*d_v, use_bias=False)
elif mode == 1:
self.qs_layers = []
self.ks_layers = []
self.vs_layers = []
for _ in range(n_head):
self.qs_layers.append(TimeDistributed(Dense(d_k, use_bias=False)))
self.ks_layers.append(TimeDistributed(Dense(d_k, use_bias=False)))
self.vs_layers.append(TimeDistributed(Dense(d_v, use_bias=False)))
self.attention = ScaledDotProductAttention()
self.w_o = TimeDistributed(Dense(d_model))
def __call__(self, q, k, v, mask=None):
d_k, d_v = self.d_k, self.d_v
n_head = self.n_head
if self.mode == 0:
qs = self.qs_layer(q) # [batch_size, len_q, n_head*d_k]
ks = self.ks_layer(k)
vs = self.vs_layer(v)
def reshape1(x):
s = tf.shape(x) # [batch_size, len_q, n_head * d_k]
x = tf.reshape(x, [s[0], s[1], n_head, s[2]//n_head])
x = tf.transpose(x, [2, 0, 1, 3])
x = tf.reshape(x, [-1, s[1], s[2]//n_head]) # [n_head * batch_size, len_q, d_k]
return x
qs = Lambda(reshape1)(qs)
ks = Lambda(reshape1)(ks)
vs = Lambda(reshape1)(vs)
if mask is not None:
mask = Lambda(lambda x:K.repeat_elements(x, n_head, 0))(mask)
head, attn = self.attention(qs, ks, vs, mask=mask)
def reshape2(x):
s = tf.shape(x) # [n_head * batch_size, len_v, d_v]
x = tf.reshape(x, [n_head, -1, s[1], s[2]])
x = tf.transpose(x, [1, 2, 0, 3])
x = tf.reshape(x, [-1, s[1], n_head*d_v]) # [batch_size, len_v, n_head * d_v]
return x
head = Lambda(reshape2)(head)
elif self.mode == 1:
heads = []; attns = []
for i in range(n_head):
qs = self.qs_layers[i](q)
ks = self.ks_layers[i](k)
vs = self.vs_layers[i](v)
head, attn = self.attention(qs, ks, vs, mask)
heads.append(head); attns.append(attn)
head = Concatenate()(heads) if n_head > 1 else heads[0]
attn = Concatenate()(attns) if n_head > 1 else attns[0]
outputs = self.w_o(head)
outputs = Dropout(self.dropout)(outputs)
return outputs, attn
Обратите внимание, что здесь все матрицы Q , K , V происходят либо из Dense() слоев, либо из TimeDistributed(Dense()) слоев.
Добавление и нормализация слоев
class LayerNormalization(Layer):
def __init__(self, eps=1e-6, **kwargs):
self.eps = eps
super(LayerNormalization, self).__init__(**kwargs)
def build(self, input_shape):
self.gamma = self.add_weight(name='gamma', shape=input_shape[-1:], initializer=Ones(), trainable=True)
self.beta = self.add_weight(name='beta', shape=input_shape[-1:], initializer=Zeros(), trainable=True)
super(LayerNormalization, self).build(input_shape)
def call(self, x):
mean = K.mean(x, axis=-1, keepdims=True)
std = K.std(x, axis=-1, keepdims=True)
return self.gamma * (x - mean) / (std + self.eps) + self.beta
def compute_output_shape(self, input_shape):
return input_shape
Слой прямой связи
Этот слой называется PositionwiseFeedForward, поскольку он учитывает положение каждого символа во входной строке:
class PositionwiseFeedForward():
def __init__(self, d_hid, d_inner_hid, dropout=0.1):
self.w_1 = Conv1D(d_inner_hid, 1, activation='relu')
self.w_2 = Conv1D(d_hid, 1)
self.layer_norm = LayerNormalization()
self.dropout = Dropout(dropout)
def __call__(self, x):
output = self.w_1(x)
output = self.w_2(output)
output = self.dropout(output)
output = Add()([output, x])
return self.layer_norm(output)
