bert,albert的快速训练和预测

  随着预训练模型越来越成熟,预训练模型也会更多的在业务中使用,本文提供了bert和albert的快速训练和部署,实际上目前的预训练模型在用起来时都大致相同。

  基于不久前发布的中文数据集chineseGLUE,将所有任务分成四大类:文本分类,句子对判断,实体识别,阅读理解。同类可以共享代码,除上面四个任务之外,还加了一个learning to rank ,基于pair wise的方式的任务,代码见:https://github.com/jiangxinyang227/bert-for-task

  具体使用见readme

  模型定义在每个项目下的model.py文件中,直接调用bert和albert的源码modeling.py将预训练模型引入,将预训练模型作为encoder部分,也可以只作为embedding层,再自己定义encoder部分,总之可以非常方便的接入下游任务网络层,尤其是当你只想使用预训练模型作为embedding层时,我们需要自己些encoder部分。

     bert_config = modeling.BertConfig.from_json_file(self.__bert_config_path)

        model = modeling.BertModel(config=bert_config,
                                   is_training=self.__is_training,
                                   input_ids=self.input_ids,
                                   input_mask=self.input_masks,
                                   token_type_ids=self.segment_ids,
                                   use_one_hot_embeddings=False)
        output_layer = model.get_pooled_output()

        hidden_size = output_layer.shape[-1].value
        if self.__is_training:
            # I.e., 0.1 dropout
            output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)

        with tf.name_scope("output"):
            output_weights = tf.get_variable(
                "output_weights", [self.__num_classes, hidden_size],
                initializer=tf.truncated_normal_initializer(stddev=0.02))

            output_bias = tf.get_variable(
                "output_bias", [self.__num_classes], initializer=tf.zeros_initializer())

            logits = tf.matmul(output_layer, output_weights, transpose_b=True)
            logits = tf.nn.bias_add(logits, output_bias)
            self.predictions = tf.argmax(logits, axis=-1, name="predictions")

  在训练时加载预训练的参数值来初始化预训练模型的变量,具体在trainer.py文件中

tvars = tf.trainable_variables()
            (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(
                tvars, self.__bert_checkpoint_path)
print("init bert model params")
tf.train.init_from_checkpoint(self.
__bert_checkpoint_path, assignment_map) print("init bert model params done") sess.run(tf.variables_initializer(tf.global_variables()))

  在预测时可以直接实例化predict.py文件中的Predictor类就会加载checkpoint模型文件,调用类中的predict方法就可以进行预测,在不需要考虑模型代码加密,模型优化等情况下,可以直接线上部署。

import json

from predict import Predictor


with open("config/tnews_config.json", "r") as fr:
    config = json.load(fr)


predictor = Predictor(config)
text = "歼20座舱盖上的两条“花纹”是什么?"
res = predictor.predict(text)
print(res)
posted @ 2019-11-18 15:15  微笑sun  阅读(1638)  评论(0编辑  收藏