Load checkpoint and finetuning using tf.estimator.

Question

We're trying to translate old training code based into a more tf.estimator.Estimator compliant code. In the initial code we fine tune an original model for a target dataset. Only some layers are loaded from the checkpoint before the training takes place using a combination of variables_to_restore and init_fn with the MonitoredTrainingSession. How can one achieve this kind of weight loading with the tf.estimator.Estimator approach ?

Answer 1

you have two options, first one is simpler:

1- use tf.train.init_from_checkpoint in your model_fn

2- model_fn returns an EstimatorSpec. You can set scaffold viaEstimatorSpec.

Answer 2

import tensorflow as tf    

def model_fn():
  # your model defintion here
  # ...

# specify your saved checkpoint path
checkpoint_path = "model.ckpt"

ws = tf.estimator.WarmStartSettings(ckpt_to_initialize_from=checkpoint_path)
est = tf.estimator.Estimator(model_fn=model_fn, warm_start_from=ws)