使用transformers和tensorflow2.0跑bert模型

1.前言

前面兩篇文章其實已經詳細介紹了bert在kaggle比賽tweet_sentiment_extraction的作用，但是該比賽是基於tensorflow2.0版本的，因此需要把代碼進行轉換。前面的兩篇文章如下鏈接：

• BERT in tweet_sentiment_extraction
• 使用BERT的兩層encoder實現tweet sentiment extraction

2. 使用tensorflow2.0 版本跑 bert模型和roberta模型

在kaggle中使用notebook參加比賽，是基於tensorflow2.0版本的，雖然不太想換版本跑，但是爲了能夠用它上面的GPU
還是借鑑了kaggle論壇上面的代碼。其中要特別感謝大佬：

• @Abhishek Thakur，其代碼啓發了我使用並行化加速訓練模型，能夠以更短的時間內訓練5-fold模型。

要在tensorflow2.0上跑bert和roberta模型，需要安裝transformers:

pip install transformers

2.1 加載transformers中的分詞包

因爲要構建bert模型的輸入，因此加載詞典，同時把輸入句子轉換成下面三個部分：

• input_ids: 把每個token轉換爲對應的id
• attention_mask： 記錄哪些詞語需要mask
• input_type_ids： 區分兩個句子，前一句子標記爲0，後一句子標記爲1

例如：

*** Example ***
tokens: happy b ##day !
input_ids: 101 3407 1038 10259 999 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
input_mask: 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
segment_ids: 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
selected_text: happy bday!

導入分詞器工具：

# set some global variables
PATH = "../input/huggingfacetransformermodels/model_classes/roberta/roberta-large-tf2-model/"
MAX_SEQUENCE_LENGTH = 128
TOKENIZER = BertWordPieceTokenizer(f"../input/my-data/vocab.txt", lowercase=True, add_special_tokens=False)

其中TOKENIZER可以直接實現分詞，轉換id等操作：

enc = TOKENIZER.encode(tweet)
input_ids_orig, offsets = enc.ids, enc.offsets

2.2 自定義bert模型層

這裏需要輸出bert的全部12層，然後取最後兩層作爲輸出。最後一層預測開始的位置start_logits，倒數第二層預測結束爲止end_logits

class BertQAModel(TFBertPreTrainedModel):
    
    DROPOUT_RATE = 0.1
    NUM_HIDDEN_STATES = 2
    
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        
        self.bert = TFBertMainLayer(config, name="bert")
        self.concat = L.Concatenate()
        self.dropout = L.Dropout(self.DROPOUT_RATE)
        self.qa_outputs = L.Dense(
            config.num_labels, 
            kernel_initializer=TruncatedNormal(stddev=config.initializer_range),
            dtype='float32',
            name="qa_outputs")
        
    @tf.function
    def call(self, inputs, **kwargs):
        # outputs: Tuple[sequence, pooled, hidden_states]
        _, _, hidden_states = self.bert(inputs, **kwargs)
        
        hidden_states = self.concat([
            hidden_states[-i] for i in range(1, self.NUM_HIDDEN_STATES+1)
        ])
        
        hidden_states = self.dropout(hidden_states, training=kwargs.get("training", False))
        logits = self.qa_outputs(hidden_states)
        start_logits, end_logits = tf.split(logits, 2, axis=-1)
        start_logits = tf.squeeze(start_logits, axis=-1)
        end_logits = tf.squeeze(end_logits, axis=-1)
        
        return start_logits, end_logits

2.3 預加載模型

利用transformers，可以快速實現預加載模型，同時transformers這個庫中已經集成了多種模型.

在加載模型之前，需要導入模型的基本設置：

config = RobertaConfig.from_json_file(os.path.join(PATH, "config.json"))
config.output_hidden_states = True
config.num_labels = 2

接下來加載模型：

model = RoBertQAModel.from_pretrained(PATH, config=config)

2.4 並行化處理（使用多線程）

本來訓練一次模型需要1100s，如果訓練兩個模型，則需要1100*2s的時間。使用多線程後，在訓練兩個模型的時候，可以把時間縮短到
1100s左右。

通過使用joblib包，來實現多線程，從而壓縮訓練時間，它的使用方法也很簡單，僅僅只需要幾行代碼就可以實現：

from joblib import Parallel, delayed
test_result = Parallel(n_jobs=num_folds, backend="threading")(delayed(run)(i) for i in range(num_folds))

• run()函數是我們自己實現的函數，裏面主要實現了模型的訓練和預測過程
• n_jobs：用來定義共有多少個線程可以實現。

實驗結果可以看出，5-fold可以縮短1000s: