Restore best checkpoint to an estimator tensorflow 2.x

Question

Briefly, I put in place a data input pipline using tensorflow Dataset API. Then, I implemented a CNN model for classification using keras, which i converted to an estimator. I feeded my estimator Train and Eval Specs with my input_fn providing input data for training and evaluation. And as final step I launched the model training with tf.estimator.train_and_evaluate

def my_input_fn(tfrecords_path):

    dataset = (...)
    return batch_fbanks, batch_labels

def build_model():
    model = tf.keras.models.Sequential()
    model.add(...)
    model.compile(...)

    return model

model = build_model()

run_config=tf.estimator.RunConfig(model_dir,save_summary_steps=100,save_checkpoints_steps=1000)
estimator = tf.keras.estimator.model_to_estimator(model,config=run_config)

def serving_input_receiver_fn():
    inputs = {'Conv1_input': tf.compat.v1.placeholder(shape=[None, 11,120,1], dtype=tf.float32)}
    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

exporter = tf.estimator.BestExporter(serving_input_receiver_fn, name="best_exporter", exports_to_keep=5)

train_spec_dnn = tf.estimator.TrainSpec(input_fn = lambda: my_input_fn(train_data_path),hooks=[hook])
eval_spec_dnn = tf.estimator.EvalSpec(input_fn = lambda: my_eval_input_fn(eval_data_path),exporters=exporter,start_delay_secs=0,throttle_secs=15)

tf.estimator.train_and_evaluate(estimator, train_spec_dnn, eval_spec_dnn)

I save the 5 best checkpoints using the tf.estimator.BestExporter as shown above. Once i finished training, i want to reload the best model and convert it to an estimator to re-evaluate the model and predict on new dataset. However my issue is in restoring the checkpoint to an estimator. I tried several solutions but each time i don't get the estimator object I need to run its evaluate and predict methods.

Just to specify more, each of the best checkpoints directory is organised as follow:

./
  variables/
      variables.data-00000-of-00002
      variables.data-00001-of-00002
      variables.index
  saved_model.pb

So the question is how can I get an estimator object from the best checkpoint so that i can use it to evaluate my model and predict on new data?

Note : I found some proposed solutions relying on TensorFlow v1 features which can not solve my problem because i work with TF v2.

Thanks a lot, any help is appreciated.

TF_Support TF_Support · Accepted Answer · 2020-04-22T10:29:35

You can use the class below created from tf.estimator.BestExporter

What it does is, except for saving the best model (.pb files and etc) it will also save the best-exported model checkpoint on a different folder.

Below is the class:

import shutil, glob, os
# import tensorflow.logging as logging
## the path where all the checkpoint reside
BEST_CHECKPOINTS_PATH_FROM = 'PATH TO ALL CHECKPOINT FILES'
## the path it will save the best exporter checkpoint files
BEST_CHECKPOINTS_PATH_TO = 'PATH TO BEST EXPORTER CHECKPOINT FILES TO BE SAVE' 

class BestCheckpointsExporter(tf.estimator.BestExporter):
      def export(self, estimator, export_path, checkpoint_path, eval_result,is_the_final_export):
          if self._best_eval_result is None or \
                    self._compare_fn(self._best_eval_result, eval_result):
                    #print('Exporting a better model ({} instead of {})...'.format(eval_result, self._best_eval_result))
              for name in glob.glob(checkpoint_path + '.*'):
                    print(name)
                    print(os.path.join(BEST_CHECKPOINTS_PATH_TO, os.path.basename(name)))
                    shutil.copy(name, os.path.join(BEST_CHECKPOINTS_PATH_TO, os.path.basename(name)))
                # also save the text file used by the estimator api to find the best checkpoint
              with open(os.path.join(BEST_CHECKPOINTS_PATH_TO, "checkpoint"), 'w') as f:
                    f.write("model_checkpoint_path: \"{}\"".format(os.path.basename(checkpoint_path)))
              self._best_eval_result = eval_result
          else:
              print('Keeping the current best model ({} instead of {}).'.format(self._best_eval_result, eval_result))

Example Usage of the Class
You will just replace the exporter by calling the class and pass the serving_input_receiver_fn.

def serving_input_receiver_fn():
    inputs = {'my_dense_input': tf.compat.v1.placeholder(shape=[None, 4], dtype=tf.float32)}
    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

exporter = BestCheckpointsExporter(serving_input_receiver_fn=serving_input_receiver_fn) 

train_spec_dnn = tf.estimator.TrainSpec(input_fn = input_fn, max_steps=5)

eval_spec_dnn = tf.estimator.EvalSpec(input_fn=input_fn,exporters=exporter,start_delay_secs=0,throttle_secs=15)
(x, y) =  tf.estimator.train_and_evaluate(keras_estimator, train_spec_dnn, eval_spec_dnn)

At this point, It will save the best-exported model checkpoint files in the folder you have specified.

For loading the checkpoint files you need to do the following steps:
Step 1: Rebuild your model instance

def build_model():
    model = tf.keras.models.Sequential()
    model.add(...)
    model.compile(...)

    return model

model = build_model()

Step 2: use the model load_weights API
Reference URL: https://www.tensorflow.org/tutorials/keras/save_and_load

ck_path = tf.train.latest_checkpoint('PATH TO BEST EXPORTER CHECKPOINT FILES')
model.load_weights(ck_path)

## From there you will be able to call the predict & evaluate the functionality of the trained model

##PREDICT
prediction = model.predict(x)

##EVALUATE
for features_batch, labels_batch in input_fn().take(1):
  model.evaluate(features_batch, labels_batch)

Note: All of these have been simulated on google colab.

Restore best checkpoint to an estimator tensorflow 2.x

1 Answers