pb模型转tfserving

摘要

训练好的pb模型转换成tfserving格式

代码

def pb_to_tfserving(infile, outfile, signature):
	"""
	infile: pb模型路径
	outfile：输出路径
	signature：签名
	"""
	input_x = "input_x" + ":0"
	output_y = "probability/Softmax" + ":0"
	with tf.gfile.GFile(infile, 'rb') as f:
		graph_def = tf.compat.v1.GraphDef()
		graph_def.ParseFromString(f.read())
	
	with tf.compat.v1.Session(graph=tf.Graph()) as sess:
		tf.import_graph_def(graph_def, name=")
		g = tf.get_default_graph()
		input_x_tensor = g.get_tensor_by_name(input_x)
		output_y_tensor = g.get_tensor_by_name(output_y)
		builder = tf.saved_model.builder.SavedModelBuilder(outfile)  ## export_path导出路径  
        signature = tf.saved_model.signature_def_utils.build_signature_def(  
            inputs={"input_x": tf.compat.v1.saved_model.utils.build_tensor_info(input_x_tensor)},
            outputs={"output": tf.compat.v1.saved_model.utils.build_tensor_info(output_y_tensor)}, method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)  
        builder.add_meta_graph_and_variables(  
            sesssess=sess, tags=[tf.saved_model.tag_constants.SERVING],  
            signature_def_map={signature_name: signature}, clear_devices=True)  ## signature_name为签名，可自定义  
        builder.save()

参考

https://developer.51cto.com/article/617105.html
https://zhuanlan.zhihu.com/p/231766484
https://zhuanlan.zhihu.com/p/477642368