def optimize_class_model(args, num_labels, logger=None): """ 加载中文分类模型 :param args: :param num_labels: :param logger: :return: """ if not logger: logger = set_logger(colored('CLASSIFICATION_MODEL, Lodding...', 'cyan'), args.verbose) try: # 如果PB文件已经存在则,返回PB文件的路径,否则将模型转化为PB文件,并且返回存储PB文件的路径 if args.model_pb_dir is None: # 获取当前的运行路径 tmp_file = os.path.join(os.getcwd(), 'predict_optimizer') if not os.path.exists(tmp_file): os.mkdir(tmp_file) else: tmp_file = args.model_pb_dir pb_file = os.path.join(tmp_file, 'classification_model.pb') if os.path.exists(pb_file): print('pb_file exits', pb_file) return pb_file import tensorflow as tf graph = tf.Graph() with graph.as_default(): with tf.Session() as sess: input_ids = tf.placeholder(tf.int32, (None, args.max_seq_len), 'input_ids') input_mask = tf.placeholder(tf.int32, (None, args.max_seq_len), 'input_mask') bert_config = modeling.BertConfig.from_json_file(os.path.join(args.bert_model_dir, 'bert_config.json')) from bert_base.train.models import create_classification_model #为了兼容多输入,增加segment_id特征,即训练代码中的input_type_ids特征。 #loss, per_example_loss, logits, probabilities = create_classification_model(bert_config=bert_config, is_training=False, #input_ids=input_ids, input_mask=input_mask, segment_ids=None, labels=None, num_labels=num_labels) segment_ids = tf.placeholder(tf.int32, (None, args.max_seq_len), 'segment_ids') loss, per_example_loss, logits, probabilities = create_classification_model(bert_config=bert_config, is_training=False, input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids, labels=None, num_labels=num_labels) # pred_ids = tf.argmax(probabilities, axis=-1, output_type=tf.int32, name='pred_ids') # pred_ids = tf.identity(pred_ids, 'pred_ids') probabilities = tf.identity(probabilities, 'pred_prob') saver = tf.train.Saver() with tf.Session() as sess: sess.run(tf.global_variables_initializer()) saver.restore(sess, tf.train.latest_checkpoint(args.model_dir)) logger.info('freeze...') from tensorflow.python.framework import graph_util tmp_g = graph_util.convert_variables_to_constants(sess, graph.as_graph_def(), ['pred_prob']) logger.info('predict cut finished !!!') # 存储二进制模型到文件中 logger.info('write graph to a tmp file: %s' % pb_file) with tf.gfile.GFile(pb_file, 'wb') as f: f.write(tmp_g.SerializeToString()) return pb_file except Exception as e: logger.error('fail to optimize the graph! %s' % e, exc_info=True)
label2id = {"0":0, "1":1} id2label = {value: key for key, value in label2id.items()} label_list = label2id.keys() num_labels = len(label_list) graph = tf.get_default_graph() with graph.as_default(): print("going to restore checkpoint") #sess.run(tf.global_variables_initializer()) input_ids_p = tf.placeholder(tf.int32, [batch_size, MAX_SEQ_LENGTH], name="input_ids") input_mask_p = tf.placeholder(tf.int32, [batch_size, MAX_SEQ_LENGTH], name="input_mask") bert_config = modeling.BertConfig.from_json_file(os.path.join(bert_dir, 'bert_config.json')) loss, per_example_loss, logits, probabilities = create_classification_model(bert_config=bert_config, is_training=False, input_ids=input_ids_p, input_mask=input_mask_p, segment_ids=None, labels=None, num_labels=num_labels) # pred_ids = tf.argmax(probabilities, axis=-1, output_type=tf.int32, name='pred_ids') # pred_ids = tf.identity(pred_ids, 'pred_ids') # probabilities = tf.identity(probabilities, 'pred_prob') saver = tf.train.Saver() # (total_loss, logits, trans, pred_ids) = create_model( # bert_config=bert_config, is_training=False, input_ids=input_ids_p, input_mask=input_mask_p, segment_ids=None, # labels=None, num_labels=num_labels, use_one_hot_embeddings=False, dropout_rate=1.0) saver = tf.train.Saver() print("model_dir: ",model_dir) saver.restore(sess, tf.train.latest_checkpoint(model_dir))
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument """The `model_fn` for TPUEstimator.""" tf.logging.info("*** Features ***") for name in sorted(features.keys()): tf.logging.info(" name = %s, shape = %s" % (name, features[name].shape)) input_ids = features["input_ids"] input_mask = features["input_mask"] segment_ids = features["segment_ids"] label_ids = features["label_ids"] is_real_example = None if "is_real_example" in features: is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32) else: is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32) is_training = (mode == tf.estimator.ModeKeys.TRAIN) (total_loss, per_example_loss, logits, probabilities) = create_classification_model(bert_config, is_training, input_ids, input_mask, segment_ids, label_ids, num_labels, use_one_hot_embeddings) tvars = tf.trainable_variables() initialized_variable_names = {} scaffold_fn = None if init_checkpoint: (assignment_map, initialized_variable_names ) = modeling.get_assignment_map_from_checkpoint( tvars, init_checkpoint) if use_tpu: def tpu_scaffold(): tf.train.init_from_checkpoint(init_checkpoint, assignment_map) return tf.train.Scaffold() scaffold_fn = tpu_scaffold else: tf.train.init_from_checkpoint(init_checkpoint, assignment_map) tf.logging.info("**** Trainable Variables ****") for var in tvars: init_string = "" if var.name in initialized_variable_names: init_string = ", *INIT_FROM_CKPT*" tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape, init_string) output_spec = None if mode == tf.estimator.ModeKeys.TRAIN: train_op = optimization.create_optimizer(total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu) output_spec = tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=total_loss, train_op=train_op, scaffold_fn=scaffold_fn) elif mode == tf.estimator.ModeKeys.EVAL: def metric_fn(per_example_loss, label_ids, logits, is_real_example): predictions = tf.argmax(logits, axis=-1, output_type=tf.int32) accuracy = tf.metrics.accuracy(labels=label_ids, predictions=predictions, weights=is_real_example) loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example) return { "eval_accuracy": accuracy, "eval_loss": loss, } eval_metrics = (metric_fn, [ per_example_loss, label_ids, logits, is_real_example ]) output_spec = tf.contrib.tpu.TPUEstimatorSpec( mode=mode, loss=total_loss, eval_metrics=eval_metrics, scaffold_fn=scaffold_fn) else: output_spec = tf.contrib.tpu.TPUEstimatorSpec( mode=mode, predictions={"probabilities": probabilities}, scaffold_fn=scaffold_fn) return output_spec