tensorflow2.0+keras简单实现BERT模型

迈不过友情╰ 2022-11-01 13:42 300阅读 0赞

### 文章目录 ###

*   *  BERT模型简介
     *  BERT模型拆解
     *   *  multi-head attention
         *  FeedForwardNetwork
         *  Transformer Encoder
         *  位置编码PositionEmbedding
         *  BERT输入Embedding
         *  Masked Language Model
         *  Next Sentence Prediction
         *  BERT base model
     *  BERT模型训练

## BERT模型简介 ##

BERT主要利用Transformer Encoder部分结合Masked Language Model，训练双向注意力模型应用到语言建模中。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2h1YW5naGFvY3M_size_16_color_FFFFFF_t_70]

## BERT模型拆解 ##

完整项目参考：[https://github.com/huanghao128/bert\_example][https_github.com_huanghao128_bert_example]

tensorflow模块导入

import numpy as np
    import tensorflow as tf
    from tensorflow import keras
    from tensorflow.keras import activations
    from tensorflow.keras import layers
    from tensorflow.keras import models
    from tensorflow.keras import backend as K

### multi-head attention ###

class MultiHeadAttention(keras.Model):
        def __init__(self, hidden_size, num_heads, **kwargs):
            super(MultiHeadAttention, self).__init__(**kwargs)
            self.hidden_size = hidden_size
            self.num_heads = num_heads
            self.head_size = hidden_size // num_heads
            self.WQ = layers.Dense(hidden_size, name="dense_query")
            self.WK = layers.Dense(hidden_size, name="dense_key")
            self.WV = layers.Dense(hidden_size, name="dense_value")
            self.dense = layers.Dense(hidden_size)
    
        def _split_heads(self, x, batch_size):
            x = tf.reshape(x, shape=[batch_size, -1, self.num_heads, self.head_size])
            return tf.transpose(x, perm=[0, 2, 1, 3]) 
    
        def call(self, query, key, value, mask):
            # query: (batch, maxlen, hidden_size)
            # key : (batch, maxlen, hidden_size)
            # value: (batch, maxlen, hidden_size)
            batch_size = tf.shape(query)[0]
    
            # shape: (batch, maxlen, hidden_size)
            query = self.WQ(query)
            key = self.WK(key)
            value = self.WV(value)
            # shape: (batch, num_heads, maxlen, head_size)
            query = self._split_heads(query, batch_size)
            key = self._split_heads(key, batch_size)
            value = self._split_heads(value, batch_size)
    
            # shape: (batch, num_heads, maxlen, maxlen)
            matmul_qk = tf.matmul(query, key, transpose_b=True)
            # 缩放 matmul_qk
            dk = tf.cast(query.shape[-1], tf.float32)
            score = matmul_qk / tf.math.sqrt(dk)
    
            if mask is not None:
                mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], dtype=tf.float32)
                score += (1 - mask) * -1e9
    
            alpha = tf.nn.softmax(score)
            context = tf.matmul(alpha, value)
            context = tf.transpose(context, perm=[0, 2, 1, 3])
            context = tf.reshape(context, (batch_size, -1, self.hidden_size))
            output = self.dense(context)
                
            return output

### FeedForwardNetwork ###

class GELU(layers.Layer):
        def __init__(self):
            super(GELU, self).__init__()
           
        def call(self, x):
            cdf = 0.5 * (1.0 + tf.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
            return x * cdf
    
    class FeedForwardNetwork(keras.Model):
        def __init__(self, dff_size, hidden_size):
            super(FeedForwardNetwork, self).__init__()
            self.dense1 = layers.Dense(dff_size)
            self.activation = GELU()
            self.dense2 = layers.Dense(hidden_size)
        
        def call(self, x):
            x = self.dense1(x)
            x = self.activation(x)
            x = self.dense2(x)
            return x

### Transformer Encoder ###

# Encoder Layer层
    class TransformerBlock(keras.Model):
        def __init__(self, hidden_size, num_heads, dff_size, rate=0.1, **kwargs):
            super(TransformerBlock, self).__init__(**kwargs)
        
            self.attention = MultiHeadAttention(hidden_size, num_heads)
            self.ffn = FeedForwardNetwork(dff_size, hidden_size)
            # Layer Normalization
            self.layernorm1 = layers.LayerNormalization(epsilon=1e-6, name="layernorm_1")
            self.layernorm2 = layers.LayerNormalization(epsilon=1e-6, name="layernorm_2")
            self.dropout1 = layers.Dropout(rate)
            self.dropout2 = layers.Dropout(rate)
            
        def call(self, x, mask, training=False):
            # multi head attention
            attn_output = self.attention(x, x, x, mask)
            attn_output = self.dropout1(attn_output, training=training)
            # residual connection
            out1 = self.layernorm1(x + attn_output)
            # ffn layer
            ffn_output = self.ffn(out1)
            ffn_output = self.dropout2(ffn_output, training=training)
            # Residual connection
            out2 = self.layernorm2(out1 + ffn_output)
            
            return out2

### 位置编码PositionEmbedding ###

**固定位置编码**

def positional_encoding(maxlen, hidden_size):
        PE = np.zeros((maxlen, hidden_size))
        for i in range(maxlen):
            for j in range(hidden_size):
                if j % 2 == 0:
                    PE[i, j] = np.sin(i / 10000 ** (j / hidden_size))
                else:
                    PE[i, j] = np.cos(i / 10000 ** ((j-1) / hidden_size))
        PE = tf.constant(PE, dtype=tf.float32)
        return PE

**可学习位置编码**

class PositionEmbedding(layers.Layer):
        def __init__(self, maxlen, hidden_size, name):
            super(PositionEmbedding, self).__init__()
            self.embedding = layers.Embedding(maxlen, hidden_size)
    
        def call(self, inputs):
            position_ids = tf.range(inputs.shape[1], dtype=tf.int32)[tf.newaxis, :]
            position_embeddings = self.embedding(position_ids)
            return position_embeddings

### BERT输入Embedding ###

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2h1YW5naGFvY3M_size_16_color_FFFFFF_t_70 1]

class BERTEmbedding(keras.Model):
        def __init__(self, vocab_size, hidden_size, maxlen, rate=0.1):
            super(BERTEmbedding, self).__init__()
            self.token_embedding = layers.Embedding(vocab_size, hidden_size, name="token_embedding")
            self.segment_embedding = layers.Embedding(2, hidden_size, name="segment_embedding")
            self.position_embedding = PositionEmbedding(maxlen, hidden_size, name="position_embedding")
            self.dropout = layers.Dropout(rate)
    
        def call(self, x, segment_ids, training=False):
            tokens_embeddings = self.token_embedding(x)
            segment_embeddings = self.segment_embedding(segment_ids)
            position_embeddings = self.position_embedding(x)
            embeddings = tokens_embeddings + segment_embeddings + position_embeddings
            embeddings = self.dropout(embeddings, training=training)
            return embeddings

### Masked Language Model ###

从一个batch的隐藏层按mask索引抽取部分数据

def gather_indexes(input_tensor, positions):
        """Gathers the vectors at the specific positions over a minibatch."""
        sequence_shape = tf.shape(input_tensor)
        batch_size = sequence_shape[0]
        seq_length = sequence_shape[1]
        width = sequence_shape[2]
        flat_offsets = tf.reshape(tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
        flat_positions = tf.reshape(positions + flat_offsets, [-1])
        flat_input_tensor = tf.reshape(input_tensor, [batch_size * seq_length, width])
        output_tensor = tf.gather(flat_input_tensor, flat_positions)
        output_tensor = tf.reshape(output_tensor, [batch_size, -1, width])
        return output_tensor

带mask的语言模型

class MaskedLanguageModel(layers.Layer):
        def __init__(self, vocab_size, name):
            super(MaskedLanguageModel, self).__init__()
            self._name = name
            b_init = tf.zeros_initializer()(shape=(vocab_size,), dtype="float32")
            self.bias = tf.Variable(initial_value=b_init, trainable=True)
            self.activation = layers.Activation(activation="softmax")
    
        def call(self, inputs, positions, weights):
            mask_inputs = gather_indexes(inputs, positions)
            out = tf.matmul(mask_inputs, weights, transpose_b=True) + self.bias
            out = self.activation(out)
            return out

### Next Sentence Prediction ###

根据输入时第一个token位置\[CLS\]表示是否为下一个句子的标示，从最后隐藏层一个位置的输出，接一个2分类的全连接层，过softmax激活做分类任务。

### BERT base model ###

class BERT(keras.Model):
        def __init__(self, vocab_size, hidden_size=768, maxlen=512, num_layers=12, nums_heads=12, rate=0.1):
            super(BERT, self).__init__()
            self.num_layers = num_layers
            self.nums_heads = nums_heads
            self.dff_size = hidden_size * 4
    
            self.embedding = BERTEmbedding(vocab_size=vocab_size, hidden_size=hidden_size, maxlen=maxlen)
            self.transformer_blocks = [TransformerBlock(hidden_size, nums_heads, self.dff_size, rate) 
                                        for i in range(num_layers)]
            
        def call(self, inputs):
            input_ids, segment_ids, input_mask, training = inputs
            # add input embedding
            x = self.embedding(input_ids, segment_ids, training)
            output_weights = self.embedding.token_embedding.weights
            # multi transformer
            for i in range(self.num_layers):
                x = self.transformer_blocks[i](x, input_mask, training)
            
            return x, output_weights

**masked\_language\_model和next sentence prediction**

class BERTMLM(keras.Model):
        def __init__(self, vocab_size, hidden_size=768, maxlen=512, num_layers=12, nums_heads=12, rate=0.1):
            super(BERTMLM, self).__init__()
    
            self.bert = BERT(vocab_size, hidden_size, maxlen, num_layers, nums_heads, rate)
    
            self.masked_language_model = MaskedLanguageModel(vocab_size=vocab_size, name="mlm")
    
            self.extract_nsp_embedding = layers.Lambda(lambda x: x[:,0], name="nsp_embedding")
            self.next_sentence = layers.Dense(2, activation='softmax', name="nsp")
    
        def call(self, inputs):
    
            input_ids, segment_ids, input_mask, masked_lm_positions, training = inputs
    
            seq_outputs, output_weights = self.bert([input_ids, segment_ids, input_mask, training])
    
            # MLM output
            mlm_output = self.masked_language_model(seq_outputs, masked_lm_positions, output_weights)
            
            # NSP output
            nsp_output = self.next_sentence(self.extract_nsp_embedding(seq_outputs))
    
            return seq_outputs, mlm_output, nsp_output

测试输出模型结构

vocab_size = 10000
    hidden_size = 768
    maxlen = 128
    mask_len = 10
    num_layers = 4
    nums_heads = 4
    
    input_ids = layers.Input(shape=(maxlen,), name="input_ids", dtype=tf.float32)
    segment_ids = layers.Input(shape=(maxlen,), name="segment_ids", dtype=tf.float32)
    input_mask = layers.Input(shape=(maxlen,), name="input_mask", dtype=tf.float32)
    masked_lm_positions = layers.Input(shape=(mask_len,), name="masked_lm_positions", dtype=tf.int32)
    masked_lm_labels = layers.Input(shape=(mask_len,vocab_size), name="masked_lm_labels", dtype=tf.int32)
    is_next_labels = layers.Input(shape=(2), name="is_next", dtype=tf.int32)
    
    bert = BERTMLM(vocab_size=vocab_size, 
                hidden_size=hidden_size, 
                maxlen=maxlen, 
                num_layers=num_layers, 
                nums_heads=nums_heads)
     
    bert_output, mlm_output, nsp_output = bert([input_ids, segment_ids, input_mask, masked_lm_positions, True])
    
    model = models.Model(inputs=[input_ids, segment_ids, input_mask, masked_lm_positions], 
                            outputs=[mlm_output, nsp_output])
    
    print(model.summary())

模型结构summary结果  
![在这里插入图片描述][20210303164611224.png]

## BERT模型训练 ##

下面是一个简单的训练bert语言模型的实例。  
**生成训练样本**

def load_dataset(data_path, tokenizer):
        all_documents = []
        with open(data_path, "r", encoding="utf8") as f:
            for line in f:
                line = line.strip()
                doc_tokens = tokenizer.tokenize(line) 
                all_documents.append(doc_tokens)
    
        max_seq_length = 128
        max_pred_per_seq = 10
        vocab_size = len(tokenizer.vocab)
    
        data_features = []
        i  = 0
        for doc_index in range(len(all_documents)):
            doc_pair_tokens = create_pair_document(all_documents, doc_index, max_seq_length, short_seq_prob=0.1)
            for tokens_a, tokens_b, is_next in doc_pair_tokens:
                example = InputExample(i, tokens_a, tokens_b, is_next, lm_labels=None)
                feature = convert_examples_to_features(example, seq_length=max_seq_length, 
                                                max_pred_per_seq=max_pred_per_seq, tokenizer=tokenizer)
                data_features.append(feature)
                i += 1
    
        input_ids, segment_ids, input_mask = [], [], []
        masked_lm_positions, masked_lm_labels, is_next_labels = [], [], []
        for feature in data_features:
            input_ids.append(feature.input_ids)
            segment_ids.append(feature.segment_ids)
            input_mask.append(feature.input_mask)
            masked_lm_positions.append(feature.masked_lm_positions)
            masked_lm_labels.append(feature.masked_lm_ids)
            is_next_labels.append(feature.is_next)
    
        masked_lm_labels = to_categorical(masked_lm_labels, num_classes=vocab_size)
        is_next_labels = to_categorical(is_next_labels, num_classes=2)
    
        input_ids = np.array(input_ids)
        segment_ids = np.array(segment_ids)
        input_mask = np.array(input_mask)
        masked_lm_positions = np.array(masked_lm_positions)
    
        data_features = [input_ids, segment_ids, input_mask, masked_lm_positions]
        data_labels = [masked_lm_labels, is_next_labels]
    
        return data_features, data_labels

**训练模型**  
model用上面已经定义好的

data_features, data_labels = load_dataset("data/sample_data.txt", tokenizer)
    print("train sample count: ", len(data_features[0]))
    
    print("start training...")
    model.fit(x=data_features, y=data_labels, validation_split=0.1, 
            batch_size=batch_size, epochs=epochs)
    
    bert_model = models.Model(inputs=model.inputs, outputs=bert_output)
    print(bert_model.summary())
    bert_model.save_weights("output/bert_model_ckpt")

完整项目参考：[https://github.com/huanghao128/bert\_example][https_github.com_huanghao128_bert_example]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2h1YW5naGFvY3M_size_16_color_FFFFFF_t_70]: /images/20221024/ca94befa514e4408a9b15a0b508837d7.png
[https_github.com_huanghao128_bert_example]: https://github.com/huanghao128/bert_example
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2h1YW5naGFvY3M_size_16_color_FFFFFF_t_70 1]: /images/20221024/609366bda0c24b0093ce5ce4cf51faaa.png
[20210303164611224.png]: /images/20221024/1620165e243e41a1b9dd61e246b7d20d.png