Skip to main content
 首页 » 编程设计

python之Tensorflow 模型保存和加载

2024年08月15日35飞鱼

如何像我们在 keras 中那样保存带有模型图的 tensorflow 模型。 我们可以保存整个模型(权重和图表)并稍后导入,而不是在预测文件中再次定义整个图表吗

在 Keras 中:

checkpoint = ModelCheckpoint('RightLane-{epoch:03d}.h5',monitor='val_loss', verbose=0,  save_best_only=False, mode='auto') 

会给出一个我们可以用来预测的h5文件

model = load_model("RightLane-030.h5") 

如何在原生 tensorflow 中做同样的事情

请您参考如下方法:

方法一:在一个文件中卡住图和权重(可能无法进行再训练)

此选项显示如何将图形和权重保存在一个文件中。它的预期用例是在训练后部署/共享模型。为此,我们将使用 protobuf (pb) 格式。

给定一个 tensorflow session (和图形),您可以生成一个 protobuf

# freeze variables 
output_graph_def = tf.graph_util.convert_variables_to_constants( 
                               sess=sess, 
                               input_graph_def =sess.graph.as_graph_def(), 
                               output_node_names=['myMode/conv/output']) 
 
# write protobuf to disk 
with tf.gfile.GFile('graph.pb', "wb") as f: 
    f.write(output_graph_def.SerializeToString()) 

其中 output_node_names 需要一个 图表结果节点的名称字符串列表(参见 tensorflow documentation)。

然后,您可以加载 protobuf 并获取带有权重的图,以便轻松地执行前向传递。

with tf.gfile.GFile(path_to_pb, "rb") as f: 
    graph_def = tf.GraphDef() 
    graph_def.ParseFromString(f.read()) 
with tf.Graph().as_default() as graph: 
    tf.import_graph_def(graph_def, name='') 
    return graph 

方法二:恢复metagraph和checkpoint(easy retraining)

如果您希望能够继续训练模型,您可能需要恢复完整的图表,即权重以及损失函数、一些梯度信息(例如对于 Adam 优化器)等。

使用时需要tensorflow生成的meta和checkpoint文件

saver = tf.train.Saver(...variables...) 
saver.save(sess, 'my-model') 

这将生成两个文件,my-model 和 my-model.meta。

从这两个文件,您可以加载图形:

  new_saver = tf.train.import_meta_graph('my-model.meta') 
  new_saver.restore(sess, 'my-model') 

更多详情可以看官方documentation .