Restore a specific checkpoint for deploying with Sagemaker and TensorFlow

Viewed 975

I'm using SageMaker for training some custom TF model I realized. During training I naturally evaluate the model multiple times in order to understand when the NN actually starts overfitting. After training I'd like to restore the model that works best (i.e. which presents minimum validation loss) and deploy it on an endpoint. However, if I use the classic Tensorflow.attach() the model that is restored corresponds with the one stored in output/model.tar.gz, which, if I got it correctly, would be the one corresponding with the last training iteration (thus it may overfits).

Is there a way for specifying to SageMaker which checkpoint restoring without necessarily retraining the model with early stopping? Even forcing SM to save in model.tar.gz the model that presents minimum validation loss and not the last one would work for me, unfortunately I didn't find any immediate way to do so...

Thank you!

1 Answers

You can use the checkpoint_path to restore the Model from a previous checkpoint:

previous_checkpoint_path = 's3://location/of/my/previous/generated/checkpoints'

tf_estimator = TensorFlow('tf-train.py', role='SageMakerRole',
                          checkpoint_path=previous_checkpoint_path
                          training_steps=10000, evaluation_steps=100,
                          train_instance_count=1, train_instance_type='ml.p2.xlarge',
                          framework_version='1.10.0')
tf_estimator.fit('s3://bucket/path/to/training/data')

For more information:

https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/tensorflow/README.rst#restoring-from-checkpoints

https://sagemaker.readthedocs.io/en/latest/sagemaker.tensorflow.html

Related