在模型上线预测时，使用pb格式模型，确定输入tensor和输出tensor，明确对应的节点即可。以下代码是最近做的ner模型的infer部分，大家可以参照修改自己的模型
import tensorflow as tf import os import pickle from bert_crf import tokenization model_dir = r'crf_output_bak/' output_graph = './pb_model/query_model.pb' bert_dir = r'chinese_L-12_H-768_A-12' # 加载label->id的词典 with open(os.path.join(model_dir, 'label2id.pkl'), 'rb') as rf: label2id = pickle.load(rf) id2label = {value: key for key, value in label2id.items()} with open(os.path.join(model_dir, 'label_list.pkl'), 'rb') as rf: label_list = pickle.load(rf) num_labels = len(label_list) tokenizer = tokenization.FullTokenizer( vocab_file=os.path.join(bert_dir, 'vocab.txt'), do_lower_case=True) def load_pb_predict(): """加载pb预测 """ text = ['深汕特别合作区'] # print('input the test sentence: {}'.format(sentence_all)) # sentence = str(input()) sentence = [[s for s in str(each)] for each in text] input_ids, input_mask, = convert(sentence) with tf.Graph().as_default(): output_graph_def = tf.GraphDef() with open(output_graph, "rb") as f: output_graph_def.ParseFromString(f.read()) tf.import_graph_def(output_graph_def, name="") res = [each.name for each in output_graph_def.node] for each in res: print(each) with tf.compat.v1.Session() as sess: sess.run(tf.compat.v1.global_variables_initializer()) t1 = time.time() input_ids_p = sess.graph.get_tensor_by_name("input_ids:0") input_mask_p = sess.graph.get_tensor_by_name("input_mask:0") # feed_dict = {input_ids_p: input_ids, input_mask_p: input_mask} # 定义输出的张量名称 output_tensor_name = sess.graph.get_tensor_by_name("viterbi/ReverseSequence_1:0") out = sess.run(output_tensor_name, feed_dict) pred_label_result = convert_id_to_label(out, id2label) t2 = time.time() print('模型预测吞吐量：{}'.format((t2-t1)/len(input_ids))) print(pred_label_result) def convert_id_to_label(pred_ids_result, idx2label): result = [] for row in range(len(pred_ids_result)): curr_seq = [] for ids in pred_ids_result[row]: if ids == 0: continue curr_label = idx2label[ids] if curr_label in ['[CLS]', '[SEP]']: continue curr_seq.append(curr_label) result.append(curr_seq) return result def convert(samples): input_ids_list = [] input_mask_list = [] for line in samples: feature = convert_single_example(0, line, label_list, 25) input_ids_list.append(feature.input_ids) input_mask_list.append(feature.input_mask) # input_ids = np.reshape([feature.input_ids],(batch_size, FLAGS.max_seq_length)) # input_mask = np.reshape([feature.input_mask],(batch_size, FLAGS.max_seq_length)) # segment_ids = np.reshape([feature.segment_ids],(batch_size, FLAGS.max_seq_length)) # label_ids =np.reshape([feature.label_ids],(batch_size, FLAGS.max_seq_length)) return input_ids_list, input_mask_list def convert_single_example(ex_index, example, label_list, max_seq_length): """ 将一个样本进行分析，然后将字转化为id, 标签转化为id,然后结构化到InputFeatures对象中 :param ex_index: index :param example: 一个样本 :param label_list: 标签列表 :param max_seq_length: :param tokenizer: :param mode: :return: """ label_map = {} # 1表示从1开始对label进行index化 for (i, label) in enumerate(label_list, 1): label_map[label] = i # 保存label->index 的map if not os.path.exists(os.path.join(model_dir, 'label2id.pkl')): with open(os.path.join(model_dir, 'label2id.pkl'), 'wb') as w: pickle.dump(label_map, w) tokens = example # tokens = .tokenize(example.text) # 序列截断 if len(tokens) >= max_seq_length - 1: tokens = tokens[0:(max_seq_length - 2)] # -2 的原因是因为序列需要加一个句首和句尾标志 ntokens = [] segment_ids = [] label_ids = [] ntokens.append("[CLS]") # 句子开始设置CLS 标志 segment_ids.append(0) # append("O") or append("[CLS]") not sure! label_ids.append("[CLS]") # O OR CLS 没有任何影响，不过我觉得O 会减少标签个数,不过拒收和句尾使用不同的标志来标注，使用LCS 也没毛病 for i, token in enumerate(tokens): ntokens.append(token) segment_ids.append(0) label_ids.append(0) ntokens.append("[SEP]") # 句尾添加[SEP] 标志 segment_ids.append(0) # append("O") or append("[SEP]") not sure! label_ids.append("[SEP]") input_ids = tokenizer.convert_tokens_to_ids(ntokens) # 将序列中的字(ntokens)转化为ID形式 input_mask = [1] * len(input_ids) # padding, 使用 while len(input_ids) < max_seq_length: input_ids.append(0) input_mask.append(0) segment_ids.append(0) # we don't concerned about it! label_ids.append(0) ntokens.append("**NULL**") # label_mask.append(0) # print(len(input_ids)) assert len(input_ids) == max_seq_length assert len(input_mask) == max_seq_length assert len(segment_ids) == max_seq_length assert len(label_ids) == max_seq_length # assert len(label_mask) == max_seq_length # 结构化为一个类 feature = InputFeatures( input_ids=input_ids, input_mask=input_mask, segment_ids=segment_ids, label_ids=label_ids, # label_mask = label_mask ) return feature class InputFeatures(object): """A single set of features of data.""" def __init__(self, input_ids, input_mask, segment_ids, label_ids, is_real_example=True): self.input_ids = input_ids self.input_mask = input_mask self.segment_ids = segment_ids self.label_ids = label_ids self.is_real_example = is_real_example if __name__ == '__main__': load_pb_predict()