如何在BERT模型上应用修剪?

时间:2020-10-20 13:04:26

标签: python tensorflow nlp bert-language-model huggingface-transformers

我已经使用ktrain(张量流包装器)训练了BERT模型来识别文本上的情绪,它虽然有效,但是却受制于缓慢的推理。这使我的模型不适用于生产环境。我已经做过一些研究,看来修剪可能会有所帮助。

Tensorflow提供了一些修剪选项,例如tf.contrib.model_pruning。问题在于它不是一种广泛使用的技术,我找不到足够简单的示例来帮助我理解如何使用它。有人可以帮忙吗?

我在下面提供我的工作代码以供参考。

import pandas as pd
import numpy as np
import preprocessor as p
import emoji
import re
import ktrain
from ktrain import text
from unidecode import unidecode
import nltk

#text preprocessing class
class TextPreprocessing:
    def __init__(self):
        p.set_options(p.OPT.MENTION, p.OPT.URL)
  
    def _punctuation(self,val): 
        val = re.sub(r'[^\w\s]',' ',val)
        val = re.sub('_', ' ',val)
        return val
  
    def _whitespace(self,val):
        return " ".join(val.split())
  
    def _removenumbers(self,val):
        val = re.sub('[0-9]+', '', val)
        return val
  
    def _remove_unicode(self, text):
        text = unidecode(text).encode("ascii")
        text = str(text, "ascii")
        return text  
    
    def _split_to_sentences(self, body_text):
        sentences = re.split(r"(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s", body_text)
        return sentences
    
    def _clean_text(self,val):
        val = val.lower()
        val = self._removenumbers(val)
        val = p.clean(val)
        val = ' '.join(self._punctuation(emoji.demojize(val)).split())
        val = self._remove_unicode(val)
        val = self._whitespace(val)
        return val
  
    def text_preprocessor(self, body_text):

        body_text_df = pd.DataFrame({"body_text": body_text},index=[1])

        sentence_split_df = body_text_df.copy()

        sentence_split_df["body_text"] = sentence_split_df["body_text"].apply(
            self._split_to_sentences)

        lst_col = "body_text"
        sentence_split_df = pd.DataFrame(
            {
                col: np.repeat(
                    sentence_split_df[col].values, sentence_split_df[lst_col].str.len(
                    )
                )
                for col in sentence_split_df.columns.drop(lst_col)
            }
        ).assign(**{lst_col: np.concatenate(sentence_split_df[lst_col].values)})[
            sentence_split_df.columns
        ]
        
        body_text_df["body_text"] = body_text_df["body_text"].apply(self._clean_text)

        final_df = (
            pd.concat([sentence_split_df, body_text_df])
            .reset_index()
            .drop(columns=["index"])
        )
        
        return final_df["body_text"]

#instantiate data preprocessing object
text1 = TextPreprocessing()

#import data
data_train = pd.read_csv('data_train_v5.csv', encoding='utf8', engine='python')
data_test = pd.read_csv('data_test_v5.csv', encoding='utf8', engine='python')

#clean the data
data_train['Text'] = data_train['Text'].apply(text1._clean_text)
data_test['Text'] = data_test['Text'].apply(text1._clean_text)

X_train = data_train.Text.tolist()
X_test = data_test.Text.tolist()

y_train = data_train.Emotion.tolist()
y_test = data_test.Emotion.tolist()

data = data_train.append(data_test, ignore_index=True)

class_names = ['joy','sadness','fear','anger','neutral']

encoding = {
    'joy': 0,
    'sadness': 1,
    'fear': 2,
    'anger': 3,
    'neutral': 4
}

# Integer values for each class
y_train = [encoding[x] for x in y_train]
y_test = [encoding[x] for x in y_test]

trn, val, preproc = text.texts_from_array(x_train=X_train, y_train=y_train,
                                                                       x_test=X_test, y_test=y_test,
                                                                       class_names=class_names,
                                                                       preprocess_mode='distilbert',
                                                                       maxlen=350)

model = text.text_classifier('distilbert', train_data=trn, preproc=preproc)

learner = ktrain.get_learner(model, train_data=trn, val_data=val, batch_size=6)

predictor = ktrain.get_predictor(learner.model, preproc)

#save the model on a file for later use
predictor.save("models/bert_model")

message = "This is a happy message"

#cleaning - takes 5ms to run
clean = text1._clean_text(message)

#prediction - takes 325 ms to run
predictor.predict_proba(clean)

1 个答案:

答案 0 :(得分:1)

火车中的distilbert模型是使用Hugging Face 变形金刚创建的,这意味着您可以使用该库来修剪模型。有关更多信息,请参见this linkthe example script。您可能需要在使用脚本之前将模型转换为PyTorch(除了对脚本本身进行一些修改之外)。该方法基于论文Are Sixteen Heads Really Better Than One?