bert模型文本分类

bert模型文本分类，实际这个东西google官方已经提供了代码，做文本分类实际是一个最简单的问题，下面用官方代码改了下，可以在低版本的tensorflow上运行，至于数据格式不再做多谈，就是input、inputmask、label，其中segment_ids可以不用做文本分类，看下代码：

import tensorflow as tf
from tensorflow.contrib.layers.python.layers import initializers
from bert import modeling
from bert import optimization
import os
import time
import datetime
from datetime import timedelta

os.environ['CUDA_VISIBLE_DEVICES'] = '0'


flags = tf.flags
FLAGS = flags.FLAGS

flags.DEFINE_string("checkpointDir",
                    "*/data/chinese_L-12_H-768_A-12/",
                    "model  save path")

bert_path = os.path.join(FLAGS.checkpointDir, "bert/")

flags.DEFINE_string(
    "bert_config_file", os.path.join(bert_path, 'bert_config.json'),
    "The config json file corresponding to the pre-trained BERT model."
)

flags.DEFINE_string(
    "init_checkpoint", os.path.join(bert_path, 'bert_model.ckpt'),
    "Initial checkpoint (usually from a pre-trained BERT model)."
)
flags.DEFINE_bool("is_training", True, "is training")

flags.DEFINE_integer("batch_size", 64, "Total batch size for training.")

flags.DEFINE_integer("tag_vocab_size", 21, "Total tag size for label")

flags.DEFINE_string("model_version", "4", "model_version ")

# 5e-5
flags.DEFINE_float("learning_rate", 5e-5, "The initial learning rate for Adam.")

flags.DEFINE_float("num_train_epochs", 3.0, "Total number of training epochs to perform.")
flags.DEFINE_float('droupout_rate', 0.5, 'Dropout rate')
flags.DEFINE_float('clip', 5, 'Gradient clip')
flags.DEFINE_float(
    "warmup_proportion", 0.1,
    "Proportion of training to perform linear learning rate warmup for. "
    "E.g., 0.1 = 10% of training.")

flags.DEFINE_integer("save_checkpoints_steps", 1000,
                     "How often to save the model checkpoint.")

flags.DEFINE_integer("iterations_per_loop", 1000,
                     "How many steps to make in each estimator call.")

flags.DEFINE_string("buckets", "resource", "buckets info")




class BertTextClassify(object):
    def __init__(self, bert_config, is_training, input_ids, input_mask,
                 segment_ids, labels, num_labels, use_one_hot_embeddings, init_checkpoint):
        self.bert_config = bert_config
        self.is_training = is_training
        self.input_ids = input_ids
        self.input_mask = input_mask
        self.segment_ids = segment_ids
        self.labels = labels
        self.num_labels = num_labels
        self.use_one_hot_embeddings = use_one_hot_embeddings
        self.init_checkpoint = init_checkpoint

        self.model = modeling.BertModel(
            config=self.bert_config,
            is_training=self.is_training,
            input_ids=self.input_ids,
            input_mask=self.input_mask,
            token_type_ids=self.segment_ids,
            use_one_hot_embeddings=self.use_one_hot_embeddings
        )
        self.output_layer=self.model.get_pooled_output()
        self.hidden_size =  self.output_layer.shape[-1].value
        self.output_weights = tf.get_variable(
            "output_weights", [ self.hidden_size,self.num_labels],
            initializer=tf.truncated_normal_initializer(stddev=0.02))
        output_bias = tf.get_variable(

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

        with tf.variable_scope("loss"):
            if is_training:
                self.output_layer = tf.nn.dropout( self.output_layer, keep_prob=0.9)

            logits = tf.matmul(self.output_layer, self.output_weights )
            self.logits = tf.nn.bias_add(logits, output_bias)
            self.probabilities = tf.nn.softmax(logits, axis=-1)
            log_probs = tf.nn.log_softmax(logits, axis=-1)
            one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

            self.per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)

            self.loss = tf.reduce_mean(self.per_example_loss)

        with tf.name_scope("train_op"):
            tvars = tf.trainable_variables()
            grads = tf.gradients(self.loss , tvars)
            global_step = tf.train.get_or_create_global_step()
            optimizer = optimization.AdamWeightDecayOptimizer(learning_rate=FLAGS.learning_rate)
            self.train_op = optimizer.apply_gradients(zip(grads, tvars), global_step)


def main(_):
    print('Loading data...')

    tag_vocab_size = FLAGS.tag_vocab_size

    input_path = FLAGS.buckets + "tfner.records*"
    files = tf.train.match_filenames_once(input_path)

    """
      inputs是你数据的输入路径
    
    """
    input_ids, input_mask, label_ids = inputs(files, batch_size=FLAGS.batch_size, num_epochs=3)
    bert_config = modeling.BertConfig.from_json_file(FLAGS.bert_config_file)
    is_training = FLAGS.is_training
    init_checkpoint = FLAGS.init_checkpoint
    use_one_hot_embeddings = False

    model = BertTextClassify(bert_config, is_training, input_ids, input_mask
                        , None, label_ids, tag_vocab_size
                        , use_one_hot_embeddings, init_checkpoint)

    tvars = tf.trainable_variables()

    if init_checkpoint:
        (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars,
                                                                                                   init_checkpoint)
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
        print("restore sucess  on cpu or gpu")

    session = tf.Session()
    session.run(tf.global_variables_initializer())
    session.run(tf.local_variables_initializer())


    print("**** Trainable Variables ****")
    for var in tvars:
        if var.name in initialized_variable_names:
            init_string = ", *INIT_FROM_CKPT*"
            print("name ={0}, shape = {1}{2}".format(var.name, var.shape,
                                                     init_string))

    print("bertlstmner  model will start train .........")

    print(session.run(files))
    saver = tf.train.Saver()
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord, sess=session)
    start_time = time.time()
    for i in range(20000):
        _, loss_train = session.run([model.train_op, model.loss])
        if i % 1000 == 0:
            end_time = time.time()
            time_dif = end_time - start_time
            time_dif = timedelta(seconds=int(round(time_dif)))
            msg = 'Iter: {0:>6}, Train Loss: {1:>6.2},' \
                  + '  Cost: {2}  Time:{3}'
            print(msg.format(i, loss_train, time_dif, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
            start_time = time.time()
        if i % 40000 == 0 and i > 0:
            saver.save(session, FLAGS.checkpointDir + "bertmodel/model.ckpt", global_step=i)
    coord.request_stop()
    coord.join(threads)
    session.close()


if __name__ == "__main__":
    tf.app.run()