tensorflow 常用代码片段
tensorflow 常用代码片段
-
加载pb文件
def load_graph(frozen_graph_filename):
with tf.io.gfile.GFile(frozen_graph_filename, “rb”) as f:
graph_def = tf.compat.v1.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def)
return graph -
写graph_def 到文件
with tf.gfile.Gfile('simplified.pb', 'wb') as fid:fid.write(graph_def.SerializeToString())
- 获取图的输入输出节点名字
def analyze_inputs_outputs(graph):ops = graph.get_operations()outputs_set = set(ops)inputs = []for op in ops:if len(op.inputs) == 0 and op.type != 'Const':inputs.append(op)else:for input_tensor in op.inputs:if input_tensor.op in outputs_set:outputs_set.remove(input_tensor.op)outputs = list(outputs_set)return (inputs, outputs)
- 简化模型，删除训练节点
graph_def = tf.graph_util.convert_variables_to_constants(sess, graph_def, [out_name])
graph_def = tf.graph_util.remove_training_nodes(graph_def)