TensorFlow Custom Estimator - 在 model_fn 中进行小的更改后恢复模型

Question

我正在使用 tf.estimator.Estimator 开发我的模型，

我写了一个 model_fn 并训练了 50,000 次迭代，现在我想对我的 model_fn 做一个小改动，例如添加一个新层。

我不想从头开始训练，我想从50,000个检查点恢复所有旧变量，从这个点继续训练。当我尝试这样做时，我得到一个 NotFoundError

tf.estimator.Estimator如何做到这一点？

Answer 1

TL;DR 从前一个检查点加载变量的最简单方法是使用函数 tf.train.init_from_checkpoint()。只需在 Estimator 的 model_fn 中调用此函数一次，就会覆盖相应变量的初始值设定项。

第一个有两个隐藏层的模型

更详细地说，假设您已经在 MNIST 上训练了第一个具有两个隐藏层的模型，名为 model_fn_1。权重保存在目录 mnist_1.

def model_fn_1(features, labels, mode):
    images = features['image']

    h1 = tf.layers.dense(images, 100, activation=tf.nn.relu, name="h1")
    h2 = tf.layers.dense(h1, 100, activation=tf.nn.relu, name="h2")

    logits = tf.layers.dense(h2, 10, name="logits")

    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    optimizer = tf.train.GradientDescentOptimizer(0.01)
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

# Estimator 1: two hidden layers
estimator_1 = tf.estimator.Estimator(model_fn_1, model_dir='mnist_1')

estimator_1.train(input_fn=train_input_fn, steps=1000)

具有三个隐藏层的第二个模型

现在我们要训练一个具有三个隐藏层的新模型 model_fn_2。我们要加载前两个隐藏层 h1 和 h2 的权重。我们使用 tf.train.init_from_checkpoint() 来做到这一点：

def model_fn_2(features, labels, mode, params):
    images = features['image']

    h1 = tf.layers.dense(images, 100, activation=tf.nn.relu, name="h1")
    h2 = tf.layers.dense(h1, 100, activation=tf.nn.relu, name="h2")
    h3 = tf.layers.dense(h2, 100, activation=tf.nn.relu, name="h3")

    assignment_map = {
        'h1/': 'h1/',
        'h2/': 'h2/'
    }
    tf.train.init_from_checkpoint('mnist_1', assignment_map)

    logits = tf.layers.dense(h3, 10, name="logits")

    loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)

    optimizer = tf.train.GradientDescentOptimizer(0.01)
    train_op = optimizer.minimize(loss, global_step=tf.train.get_global_step())

    return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)

# Estimator 2: three hidden layers
estimator_2 = tf.estimator.Estimator(model_fn_2, model_dir='mnist_2')

estimator_2.train(input_fn=train_input_fn, steps=1000)

assignment_map 会将检查点中作用域 h1/ 中的每个变量加载到新作用域 h1/ 中，与 h2/ 相同。不要忘记最后的 / 让 TensorFlow 知道它是一个变量范围。

我找不到使用预制估算器来完成这项工作的方法，因为您无法更改它们 model_fn。

TensorFlow Custom Estimator - 在 model_fn 中进行小的更改后恢复模型

TensorFlow Custom Estimator - Restore model after small changes in model_fn

tensorflow

tensorflow-datasets

tensorflow-estimator

第一个有两个隐藏层的模型

具有三个隐藏层的第二个模型