我正在使用model_fn
来开发我的模型,
我写了model_fn
并训练了50,000次迭代,现在我想在我的NotFoundError
中进行一些小改动,例如添加一个新图层。
我不想从头开始训练,我想从50,000检查点恢复所有旧变量,并从这一点继续训练。当我尝试这样做时,我得到tf.estimator.Estimator
如何使用{{1}}完成此操作?
答案 0 :(得分:7)
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
。