在@Vishnuvardhan Janapati的回答之后，这里是另一种在TensorFlow 2.0.0下保存和重载自定义层/度量/损失模型的方法

import tensorflow as tf
from tensorflow.keras.layers import Layer
from tensorflow.keras.utils.generic_utils import get_custom_objects

# custom loss (for example)  
def custom_loss(y_true,y_pred):
  return tf.reduce_mean(y_true - y_pred)
get_custom_objects().update({'custom_loss': custom_loss}) 

# custom loss (for example) 
class CustomLayer(Layer):
  def __init__(self, ...):
      ...
  # define custom layer and all necessary custom operations inside custom layer

get_custom_objects().update({'CustomLayer': CustomLayer})  

通过这种方式，一旦您执行了这些代码，并使用tf.keras.models保存了您的模型。Save_model或model。save或ModelCheckpoint回调，您可以重新加载您的模型，而不需要精确的自定义对象，就像这样简单

new_model = tf.keras.models.load_model("./model.h5"})

这里所有的答案都很棒，但我想补充两点。

首先，详细说明@user7505159的答案，“。添加到要恢复的文件名的开头可能很重要。

例如，您可以保存没有“的图形。/"在文件名中如下所示:

# Some graph defined up here with specific names

saver = tf.train.Saver()
save_file = 'model.ckpt'

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.save(sess, save_file)

但是为了恢复图形，您可能需要在前面加上一个"。/"到file_name:

# Same graph defined up here

saver = tf.train.Saver()
save_file = './' + 'model.ckpt' # String addition used for emphasis

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    saver.restore(sess, save_file)

你并不总是需要“。/”，但是它会根据你的环境和TensorFlow版本而导致问题。

它还想提到sess.run(tf.global_variables_initializer())在恢复会话之前可能很重要。

如果在尝试恢复保存的会话时收到关于未初始化变量的错误，请确保在保存程序之前包含sess.run(tf.global_variables_initializer())。恢复(sess, save_file)行。这样你就不用头疼了。

Tensorflow 2.6:它现在变得更简单了，你可以用两种格式保存模型

Saved_model (tf服务兼容) H5或HDF5

以两种格式保存模型:

 from tensorflow.keras import Model
 inputs = tf.keras.Input(shape=(224,224,3))
 y = tf.keras.layers.Conv2D(24, 3, activation='relu', input_shape=input_shape[1:])(inputs)
 outputs = tf.keras.layers.Dense(5, activation=tf.nn.softmax)(y)
 model = tf.keras.Model(inputs=inputs, outputs=outputs)
 model.save("saved_model/my_model") #To Save in Saved_model format
 model.save("my_model.h5") #To save model in H5 or HDF5 format

以两种格式加载模型

import tensorflow as tf
h5_model = tf.keras.models.load_model("my_model.h5") # loading model in h5 format
h5_model.summary()
saved_m = tf.keras.models.load_model("saved_model/my_model") #loading model in saved_model format
saved_m.summary()

对于TensorFlow版本< 0.11.0RC1:

保存的检查点包含模型中的变量值，而不是模型/图本身，这意味着当您恢复检查点时，图应该是相同的。

这里有一个线性回归的例子，其中有一个训练循环，保存变量检查点，还有一个评估部分，将恢复之前运行中保存的变量并计算预测。当然，如果你愿意，你也可以恢复变量并继续训练。

x = tf.placeholder(tf.float32)
y = tf.placeholder(tf.float32)

w = tf.Variable(tf.zeros([1, 1], dtype=tf.float32))
b = tf.Variable(tf.ones([1, 1], dtype=tf.float32))
y_hat = tf.add(b, tf.matmul(x, w))

...more setup for optimization and what not...

saver = tf.train.Saver()  # defaults to saving all variables - in this case w and b

with tf.Session() as sess:
    sess.run(tf.initialize_all_variables())
    if FLAGS.train:
        for i in xrange(FLAGS.training_steps):
            ...training loop...
            if (i + 1) % FLAGS.checkpoint_steps == 0:
                saver.save(sess, FLAGS.checkpoint_dir + 'model.ckpt',
                           global_step=i+1)
    else:
        # Here's where you're restoring the variables w and b.
        # Note that the graph is exactly as it was when the variables were
        # saved in a prior training run.
        ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            saver.restore(sess, ckpt.model_checkpoint_path)
        else:
            ...no checkpoint found...

        # Now you can run the model to get predictions
        batch_x = ...load some data...
        predictions = sess.run(y_hat, feed_dict={x: batch_x})

下面是变量文档，涵盖了保存和恢复。这是保存程序的文档。

模型有两个部分，模型定义，由Supervisor保存为图。模型目录中的PBTXT和张量的数值，保存到检查点文件，如model.ckpt-1003418。

可以使用tf恢复模型定义。import_graph_def，并且使用Saver恢复权重。

然而，Saver使用特殊的集合保存附加到模型Graph的变量列表，并且这个集合没有使用import_graph_def初始化，所以您目前不能同时使用这两者(这在我们的路线图中进行修复)。现在，您必须使用Ryan Sepassi的方法——手动构造具有相同节点名称的图，并使用Saver将权重加载到其中。

(或者，您可以通过使用import_graph_def，手动创建变量和使用tf.add_to_collection(tf.GraphKeys. collection)来破解它。变量，变量)为每个变量，然后使用Saver)

在大多数情况下，使用tf.train.Saver从磁盘保存和恢复是最好的选择:

... # build your model
saver = tf.train.Saver()

with tf.Session() as sess:
    ... # train the model
    saver.save(sess, "/tmp/my_great_model")

with tf.Session() as sess:
    saver.restore(sess, "/tmp/my_great_model")
    ... # use the model

您还可以保存/恢复图结构本身(详细信息请参阅MetaGraph文档)。默认情况下，保存程序将图形结构保存到.meta文件中。您可以调用import_meta_graph()来恢复它。它恢复图形结构并返回一个你可以用来恢复模型状态的保护程序:

saver = tf.train.import_meta_graph("/tmp/my_great_model.meta")

with tf.Session() as sess:
    saver.restore(sess, "/tmp/my_great_model")
    ... # use the model

然而，在某些情况下，您需要更快的方法。例如，如果您实现了早期停止，那么您希望在训练期间每次模型改进时都保存检查点(在验证集上测量)，然后如果一段时间内没有进展，则希望回滚到最佳模型。如果每次模型改进时都将其保存到磁盘，则会极大地降低训练速度。诀窍是将变量状态保存到内存中，然后稍后恢复它们:

... # build your model

# get a handle on the graph nodes we need to save/restore the model
graph = tf.get_default_graph()
gvars = graph.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
assign_ops = [graph.get_operation_by_name(v.op.name + "/Assign") for v in gvars]
init_values = [assign_op.inputs[1] for assign_op in assign_ops]

with tf.Session() as sess:
    ... # train the model

    # when needed, save the model state to memory
    gvars_state = sess.run(gvars)

    # when needed, restore the model state
    feed_dict = {init_value: val
                 for init_value, val in zip(init_values, gvars_state)}
    sess.run(assign_ops, feed_dict=feed_dict)

A quick explanation: when you create a variable X, TensorFlow automatically creates an assignment operation X/Assign to set the variable's initial value. Instead of creating placeholders and extra assignment ops (which would just make the graph messy), we just use these existing assignment ops. The first input of each assignment op is a reference to the variable it is supposed to initialize, and the second input (assign_op.inputs[1]) is the initial value. So in order to set any value we want (instead of the initial value), we need to use a feed_dict and replace the initial value. Yes, TensorFlow lets you feed a value for any op, not just for placeholders, so this works fine.