Transformer示例【tensorflow 2.2】

最新推荐文章于 2024-07-24 16:14:10 发布

gdhy9064

最新推荐文章于 2024-07-24 16:14:10 发布

阅读量4.2k

点赞数 3

分类专栏： NLP 文章标签： tensorflow nlp

本文链接：https://blog.csdn.net/gdhy9064/article/details/106944842

版权

NLP 专栏收录该内容

1 篇文章 0 订阅

订阅专栏

文章目录

前言
数据准备
模型构建
训练模型
测试模型
结语

前言

Transformer是NLP中一个比较基础的模型，鼎鼎大名的BERT正是基于它而来的，因此了解它的原理及实现对NLP的学习很有必要。关于Transformer的原理，可参考https://zhuanlan.zhihu.com/p/44121378。本文的示例模型参考了tensorflow官方教程，模型代码基于tensorflow已有的API进行了重构，数据集预处理及测试完全照搬教程的代码。

数据准备

教程中所用数据是”葡萄牙语-英语翻译数据集“，其可使用tfds导入，导入代码如下

examples = tfds.load('ted_hrlr_translate/pt_to_en', as_supervised=True)
train_examples, val_examples = examples['train'], examples['validation']

接下来从训练数据集创建自定义子词分词器

tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (en.numpy() for pt, en in train_examples),
    target_vocab_size=2**13
)
tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (pt.numpy() for pt, en in train_examples),
    target_vocab_size=2**13
)

sample_string = 'Transformer is awesome.'
tokenized_string = tokenizer_en.encode(sample_string)
print ('Tokenized string is {}'.format(tokenized_string))

original_string = tokenizer_en.decode(tokenized_string)
print ('The original string: {}'.format(original_string))
assert original_string == sample_string

这里的分词器以词根作为划分的基础，如playing会分为play和ing分别编码，如果遇到生词则将单词分解，如Transformer分为T，ran，former后进行编码。并且，这里分词器编码与解码过程是可逆的。
之后开始构造用于训练的数据集

BUFFER_SIZE = 20000
BATCH_SIZE = 64
MAX_LENGTH = 40

def encode(lang1, lang2):
    lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(lang1.numpy()) + [tokenizer_pt.vocab_size+1]
    lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(lang2.numpy()) + [tokenizer_en.vocab_size+1]
    return lang1, lang2

def tf_encode(pt, en):
    result_pt, result_en = tf.py_function(encode, [pt, en], [tf.int64, tf.int64])
    result_pt.set_shape([None])
    result_en.set_shape([None])
    return result_pt, result_en

def filter_max_length(x, y, max_length=MAX_LENGTH):
    return tf.logical_and(tf.size(x) <= max_length,
                        tf.size(y) <= max_length)

train_dataset = train_examples.map(tf_encode)
train_dataset = train_dataset.filter(filter_max_length)
# 将数据集缓存到内存中以加快读取速度。
train_dataset = train_dataset.cache()
train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(BATCH_SIZE)
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

这里encode函数对数据进行预处理，它将tokenizer_en.vocab_size及tokenizer_en.vocab_size+1分别作为句子起始及结束的编码，加到每一句的编码上，对葡萄牙语的处理同理。这里还设置数据集的批次大小为64，并且只选取编码长度在40以内的句子作为训练数据。

模型构建

位置编码层

位置编码层的作用是根据词向量输入x与指定的词向量维度 $d_{model}$ （这里要求为偶数），计算位置编码并加到x上。位置编码的计算公式如下
$2i)=\sin\frac{pos}{10000^\frac{2i}{d_{model}}}$

$2i+1)=\cos\frac{pos}{10000^\frac{2i}{d_{model}}}$
其中 $i$ 代表词向量分量下标（ $0\leq i<d_{model}$ ）， $p o s$ 代表词向量处于句子中的位置（ $0\leq pos<$ 句子编码的长度）。
位置编码层的代码如下

class PositionalLayer(Layer):
    def __init__(self, input_units):
        super().__init__()
        assert input_units % 2 == 0, "Input_units should be even."
        self.base = K.constant((1 / 10000) ** (np.arange(input_units / 2) * 2 / input_units))

    def call(self, x):
        length = K.shape(x)[1]
        angles = K.transpose(K.tile(self.base[:, None], [1, length]) * K.arange(0, length, dtype='float32'))
        positional_encoding = K.concatenate([K.sin(angles), K.cos(angles)], axis=1)
        return x + positional_encoding

这里构造函数参数input_units为 $d_{model}$ , 在位置编码层初始化时，先计算固定部分self.base= $\frac{1}{10000^\frac{2i}{d_{model}}}$ 。之后根据输入分别计算sin及cos部分，原本的位置编码间隔放置sin及cos部分，这里直接将位置编码前半部分置为sin部分，后半部分置为cos部分，这样做是等效的。

多头注意力

多头注意力以自注意力为基础，tensorflow里面提供了注意力层，其API¹如下

tf.keras.layers.Attention(
    use_scale=False, **kwargs
)

其中use_scale为True代表使用对点积进行缩放，**kwargs还接受两个参数causal和dropout，causal为True等效与教程中使用前瞻遮挡，其含义是计算当前编码向量 $q_{pos}$ 的注意力输出时只使用位置 $p o s$ 之前的数据 $k_{j}, v_j(0\leq j\leq pos)$ ，dropout（ $0\leq dropout \leq 1$ )的作用是对注意力比例作dropout，与教程里面对注意力输出作dropout有区别，后面会用它代替一下。
注意力层的call函数可接受三个参数inputs，mask，training，inputs可为列表[q, v]或[q, v, k], mask为列表[q_mask, v_mask], mask（False表示遮挡）用于遮挡q和v对应位置的编码向量，这里用于遮挡训练数据的padding部分。
多头注意力由多个注意力层组成，对每个注意力层来说，首先对q, k, v作线性变换，输入到注意力层得到输出，最后将各个注意力层的输出合并。传统做法是按这种思路设置多个注意力层及对应的几个全连接层，使用for循环计算，只是这种做法比较慢。我们可以将这些全连接层连接起来得到三个对应q, k, v的全连接层 $dense_q$ , $dense_k$ , $dense_v$ ，计算出结果后再将结果的向量维度按注意力头数分割，最后只用一个注意力层计算出全部结果。如注意力头数为2，那么有 $dense_{qi}$ , $dense_{ki}$ , $dense_{vi}$ （ $0\leq i<2$ )，假设units都为32，那么合并后 $dense_q$ , $dense_k$ , $dense_v$ 的units都为64，通过变换后输出的维度中前32为对应第一个注意力层前的线性变换结果，后32为对应第二个注意力层前的线性变换结果，代码如下

class MultiHeadAttention(Layer):
    def __init__(self, input_units, head_units, transform_units, **kargs):
        super().__init__()
        self.head_units = head_units    
        self.dense_q = TimeDistributed(Dense(transform_units * head_units))
        self.dense_k = TimeDistributed(Dense(transform_units * head_units))
        self.dense_v = TimeDistributed(Dense(transform_units * head_units))
        self.attention = Attention(**kargs)   
        self.dense_output = TimeDistributed(Dense(input_units))


    def _split_and_concat(self, x):
        return K.concatenate(tf.split(x, self.head_units, axis=-1), axis=0)


    def call(self, q, v, q_mask, v_mask):
        k = v
        q_transform = self._split_and_concat(self.dense_q(q))
        v_transform = self._split_and_concat(self.dense_v(v))
        k_transform = self._split_and_concat(self.dense_k(k))
        
        head_concat = K.concatenate(
            tf.split(
                self.attention(
                    [q_transform, v_transform, k_transform],
                    mask=[
                        K.tile(q_mask, [self.head_units, 1]),
                        K.tile(v_mask, [self.head_units, 1])
                    ]
                ),
                self.head_units, 
                axis=0
            ),
            axis=-1
        )
        return self.dense_output(head_concat)

其中构造函数参数input_units为 $d_{model}$ ， head_units为注意力头数，transform_units为原本每个注意力层对应的全连接层的units，**kargs用于接收注意力层的其他参数。

残差归一化层

这一层的作用是包装其他层，将输入x输入到所包装的层后获取输出y，对x + y进行层归一化，代码如下

class ResNorm(Layer):
    def __init__(self, sequential):
        super().__init__()
        self.sequential = sequential
        self.layer_norm = LayerNormalization()

    def call(self, x):
        return self.layer_norm(x + self.sequential(x))

    
class ResNormAttention(Layer):
    def __init__(self, attention_layer):
        super().__init__()
        self.attention_layer = attention_layer
        self.layer_norm = LayerNormalization()

    def call(self, q, v, q_mask, v_mask):
        return self.layer_norm(q + self.attention_layer(q, v, q_mask, v_mask))

这里由于注意力层的输入有多个，特别设置一层用来包装注意力层。

编码器层

编码器层由经过残差归一化层包装的多头注意力及后面的同样经残差归一化层包装的前向反馈层组成，代码如下

class EncoderLayer(Layer):
    def __init__(self, input_units, head_units, transform_units, dropout, ffn_units):
        super().__init__()
        self.attention = ResNormAttention(
            MultiHeadAttention(
                input_units, 
                head_units,
                transform_units, 
                use_scale=True, 
                dropout=dropout
            )
        )
        self.ffn = ResNorm(Sequential([
            TimeDistributed(Dense(ffn_units, activation='relu')),
            TimeDistributed(Dense(input_units)),
        ]))
        

    def call(self, encoding, padding):
        return self.ffn(self.attention(encoding, encoding, padding, padding))

编码器

编码器由词嵌入层、位置编码层及多个编码器层组成，这里按照教程的做法给词嵌入层输出乘上 $\sqrt {d_{model}}$ , 教程中这里还有dropout层，这里没有加上。代码如下


class Encoder(Layer):
    def __init__(self, embedding_input_dim, embedding_output_dim, layer_units, head_units, transform_units, dropout, ffn_units):
        super().__init__()
        self.embedding_output_dim = embedding_output_dim
        self.embedding_layer = Embedding(embedding_input_dim, embedding_output_dim)
        self.pos_layer = PositionalLayer(embedding_output_dim)
        self.encoder_layers = [EncoderLayer(embedding_output_dim, head_units, transform_units, dropout, ffn_units) for _ in range(layer_units)]
    
    def call(self, embedding_input, padding):
        encoding = self.embedding_layer(embedding_input) * K.sqrt(K.constant(self.embedding_output_dim))
        encoding = self.pos_layer(encoding)
        for layer in self.encoder_layers:
            encoding = layer(encoding, padding)
        return encoding

其中构造函数参数embedding_input_dim为包含起始与终止编码的分词器大小，embedding_output_dim为词向量维数，layer_units为编码器层层数，其他参数同上。

解码器层

解码器层由两层经过残差归一化层包装的多头注意力及后面的同样经残差归一化层包装的前向反馈层组成，第一层多头注意力构造时置参数causal为True, 它的q和v都为待编码向量decoding，第二层多头注意力的q为第一层经过残差归一化多头注意力的输出，v为编码器的输出，代码如下

class DecoderLayer(Layer):
    def __init__(self, input_units, head_units, transform_units, dropout, ffn_units):
        super().__init__()
        self.attention1 = ResNormAttention(
            MultiHeadAttention(
                input_units, 
                head_units,
                transform_units, 
                use_scale=True,
                causal=True,
                dropout=dropout
            )
        ) 
        self.attention2 = ResNormAttention(
            MultiHeadAttention(
                input_units, 
                head_units,
                transform_units, 
                use_scale=True,
                dropout=dropout
            )
        )
        self.ffn = ResNorm(Sequential([
            TimeDistributed(Dense(ffn_units, activation='relu')),
            TimeDistributed(Dense(input_units)),
        ]))


    def call(self, encoding, decoding, encoding_padding, decoding_padding):
        return self.ffn(
            self.attention2(
                self.attention1(decoding, decoding, decoding_padding, decoding_padding),
                encoding,
                decoding_padding,
                encoding_padding
            )
        )

解码器

解码器层由词嵌入层、位置编码层、多个解码器层及最后一层变换组成，这里按照教程的做法给词嵌入层输出乘上 $\sqrt {d_{model}}$ , 解码器最终输出为维数为包含起始与终止编码的分词器大小，没有经过softmax的解码向量，教程中这里还有dropout层，这里没有加上。代码如下

class Decoder(Layer):
    def __init__(self, embedding_input_dim, embedding_output_dim, layer_units, head_units, transform_units, dropout, ffn_units):
        super().__init__()
        self.embedding_output_dim = embedding_output_dim
        self.embedding_layer = Embedding(embedding_input_dim, embedding_output_dim)
        self.pos_layer = PositionalLayer(embedding_output_dim)
        self.decoder_layers = [DecoderLayer(embedding_output_dim, head_units, transform_units, dropout, ffn_units) for _ in range(layer_units)]
        self.final_layer = TimeDistributed(Dense(embedding_input_dim))
        
    def call(self, encoding, embedding_input, encoding_padding, decoding_padding):
        decoding = self.embedding_layer(embedding_input) * K.sqrt(K.constant(self.embedding_output_dim))
        decoding = self.pos_layer(decoding)
        for layer in self.decoder_layers:
            decoding = layer(encoding, decoding, encoding_padding, decoding_padding)
        decoding = self.final_layer(decoding)
        return decoding

这里的参数含义同编码器。

Transformer

Transformer由编码器和解码器组成，通过自定义优化器学习率策略，使用经过padding的交叉熵损失训练，训练时的解码器输出的每个位置是在假设已知当前位置之前每个位置的情况下得到的下一位置预测编码，将它与实际编码一起计算损失，代码如下

ENCODER_EMBEDDING_INPUT_DIM = tokenizer_pt.vocab_size + 2
ENCODER_EMBEDDING_OUTPUT_DIM = 128
DECODER_EMBEDDING_INPUT_DIM = tokenizer_en.vocab_size + 2 
DECODER_EMBEDDING_OUTPUT_DIM = 128
LAYER_UNITS = 4
HEAD_UNITS = 8
TRANSFORM_UNITS = ENCODER_EMBEDDING_OUTPUT_DIM // HEAD_UNITS
FFN_UNITS = 512
DROPOUT = 0.1

encoder = Encoder(
    ENCODER_EMBEDDING_INPUT_DIM,
    ENCODER_EMBEDDING_OUTPUT_DIM,
    LAYER_UNITS,
    HEAD_UNITS,
    TRANSFORM_UNITS,
    DROPOUT,
    FFN_UNITS
)

decoder = Decoder(
    DECODER_EMBEDDING_INPUT_DIM,
    DECODER_EMBEDDING_OUTPUT_DIM,
    LAYER_UNITS,
    HEAD_UNITS,
    TRANSFORM_UNITS,
    DROPOUT,
    FFN_UNITS
)

def loss_func(decoding_real, decoding_pred):
    mask = K.not_equal(decoding_real, 0)
    # from_logits=True表示预测的解码向量没有经过softmax
    loss = tf.keras.losses.sparse_categorical_crossentropy(decoding_real, decoding_pred, from_logits=True)
    mask = tf.cast(mask, dtype=loss.dtype)
    loss *= mask
    return K.mean(loss)



encoder_embedding_input = Input([None], dtype='int64')
decoder_embedding_input = Input([None], dtype='int64')
# 遮挡编码为0的位置，编码0在分词器中为空串，不会出现在句子中间
encoding_padding = K.not_equal(encoder_embedding_input, 0)
decoding_padding = K.not_equal(decoder_embedding_input, 0)
encoding = encoder(encoder_embedding_input, encoding_padding)
decoding = decoder(encoding, decoder_embedding_input, encoding_padding, decoding_padding)


transformer = Model(
    inputs=[
        encoder_embedding_input,
        decoder_embedding_input
    ],
    outputs=decoding
)

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, embedding_dim, warmup_steps=4000):
        super(CustomSchedule, self).__init__()
        self.embedding_dim = tf.cast(embedding_dim, tf.float32)
        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps ** -1.5)
        return tf.math.rsqrt(self.embedding_dim) * tf.math.minimum(arg1, arg2)

learning_rate = CustomSchedule(ENCODER_EMBEDDING_OUTPUT_DIM)

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)
transformer.compile(optimizer=optimizer, loss=loss_func)
print(transformer.summary())

Transformer的模型结构输出如下
在这里插入图片描述训练过程的学习率将随训练步数的变化如下

训练模型

使用已预处理的训练数据，设置训练轮数为20，训练代码如下

EPOCHS = 20
for epoch in range(EPOCHS):
    for batch, (inp, tar) in enumerate(train_dataset):
        print(f'{epoch}-{batch}')
        tar_inp = tar[:, :-1]
        tar_real = tar[:, 1:]
        transformer.fit([inp, tar_inp], tar_real, batch_size=BATCH_SIZE)

测试模型

首先实现对编码输入进行翻译的函数translate，不同于训练时编码器直接输出结果，翻译时需要像自回归模型一个一个地预测之后的单词，代码如下

def translate(sentence):
    start_token = [tokenizer_pt.vocab_size]
    end_token = [tokenizer_pt.vocab_size + 1]

    # 输入语句是葡萄牙语，增加开始和结束标记
    sentence = start_token + tokenizer_pt.encode(sentence) + end_token
    encoder_input = K.expand_dims(sentence, 0)

    # 因为目标是英语，输入 transformer 的第一个词应该是
    # 英语的开始标记。
    decoder_input = [tokenizer_en.vocab_size]
    output = K.expand_dims(decoder_input, 0)
    encoding = encoder(encoder_input, encoder_input != 0)
    for i in range(MAX_LENGTH):
        # predictions.shape == (batch_size, seq_len, vocab_size)
        predictions = decoder(encoding, output, encoder_input != 0, output != 0)
        # 从 seq_len 维度选择最后一个词
        predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)
        predicted_id = K.cast(np.argmax(predictions, axis=-1), tf.int32)    
        # 如果 predicted_id 等于结束标记，就返回结果
        if predicted_id == tokenizer_en.vocab_size + 1:
            break 
        # 连接 predicted_id 与输出，作为解码器的输入传递到解码器。
        output = K.concatenate([output, predicted_id], axis=-1)
    return tokenizer_en.decode([i for i in K.squeeze(output, axis=0) if i < tokenizer_en.vocab_size])

这个过程中前面已经计算得到的预测编码会重复计算，但实际上每次我们只需要最后一个编码，在当前模型结构下无法解决重复计算的问题。
之后利用教程的样例进行测试，测试代码如下

print('Predicted translation: ' + translate("este é um problema que temos que resolver."))
print("Real translation: this is a problem we have to solve .")

print('Predicted translation: ' + translate("os meus vizinhos ouviram sobre esta ideia."))
print("Real translation: and my neighboring homes heard about this idea .")

print('Predicted translation: ' + translate("os meus vizinhos ouviram sobre esta ideia."))
print("Real translation: and my neighboring homes heard about this idea .")

测试结果如下
在这里插入图片描述

结语

至此我们把Transformer的结构过了一遍，相信这对于理解Transformer有不小的帮助，完整代码在此。

Attention API参考 https://tensorflow.google.cn/api_docs/python/tf/keras/layers/Attention ↩︎

gdhy9064

关注

3
点赞
踩
41

收藏

觉得还不错? 一键收藏
2
评论
Transformer示例【tensorflow 2.2】

本文利用tensorflow已有api实现了一个简易的Transformer，易于读者通过代码理解Transformer的思想，其中数据集预处理及测试照搬官方教程的代码。
复制链接

扫一扫

专栏目录