def serving_input_fn():
label_ids = tf.placeholder(tf.int32, [None, 3], name='label_ids')
input_ids = tf.placeholder(tf.int32, [None, 200], name='input_ids')
input_mask = tf.placeholder(tf.int32, [None, 200], name='input_mask')
segment_ids = tf.placeholder(tf.int32, [None, 200], name='segment_ids')
input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
'label_ids': label_ids,
'input_ids': input_ids,
'input_mask': input_mask,
'segment_ids': segment_ids,
})()
return input_fn
在eval阶段生成pb格式，eval时加入如下代码：
if FLAGS.do_eval:
estimator._export_to_tpu = False
estimator.export_savedmodel("pbmodels", serving_input_fn)
