保存tensorflow训练好的模型
时间: 2023-10-08 16:06:23 浏览: 48
可以使用 TensorFlow 提供的 `tf.train.Saver` 对象,将训练好的模型保存为一个 checkpoint 文件或者一个 pb 文件,以便后续使用。例如,下面的代码展示了如何将模型保存为一个 checkpoint 文件:
```
saver = tf.train.Saver()
saver.save(sess, 'model.ckpt')
```
或者将模型保存为一个 pb 文件:
```
input_tensor = ...
output_tensor = ...
signature = tf.saved_model.signature_def_utils.build_signature_def(
inputs={'input': tf.saved_model.utils.build_tensor_info(input_tensor)},
outputs={'output': tf.saved_model.utils.build_tensor_info(output_tensor)},
method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME)
builder = tf.saved_model.builder.SavedModelBuilder('exported_model')
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
signature,
})
builder.save()
```
这里假设 `input_tensor` 和 `output_tensor` 分别表示模型的输入和输出_tensors。保存为 checkpoint 文件时,`sess` 表示当前 TensorFlow 的会话对象,而保存为 pb 文件时,`sess` 则不需要进行任何计算,只需要构建好合适的 `input_tensor` 和 `output_tensor`,并且创建一个 `signature_def` 对象即可。保存 pb 文件时,还需要设置一个 `tag_constants.SERVING` 的标记,表示这个模型用于 TensorFlow Serving。