Amazon SageMaker
Developer Guide

Use TensorFlow with Amazon SageMaker

You can use Amazon SageMaker to train and deploy a model using custom TensorFlow code. The Amazon SageMaker Python SDK TensorFlow estimators and models and the Amazon SageMaker open-source TensorFlow containers make writing a TensorFlow script and running it in Amazon SageMaker easier.

The Amazon SageMaker Python SDK supports two formats for TensorFlow training scripts. For TensorFlow versions 1.11 and later, you can use script mode TensorFlow training scripts. For TensorFlow versions 1.12 and earlier, you can use legacy mode TensorFlow training scripts.

For API reference for the Amazon SageMaker Python SDK TensorFlow classes, see https://sagemaker.readthedocs.io/en/stable/sagemaker.tensorflow.html.

For information about how to build and contribute to the Amazon SageMaker TensorFlow container, see the GitHub repository at https://github.com/aws/sagemaker-tensorflow-container.

Use TensorFlow Script Mode

For TensorFlow versions 1.11 and later, the Amazon SageMaker Python SDK supports script mode training scripts. Script mode has the following advantages over legacy mode training scripts:

  • Script mode training scripts are more similar to training scripts you write for TensorFlow in general, so it is easier to modify your existing TensorFlow training scripts to work with Amazon SageMaker.

  • Script mode supports both Python 2.7- andPython 3.6-compatible source files.

  • Script mode supports Horovod for distributed training.

For information about writing TensorFlow script mode training scripts and using TensorFlow script mode estimators and models with Amazon SageMaker, see https://sagemaker.readthedocs.io/en/stable/using_tf.html.

For information about TensorFlow versions supported by the Amazon SageMaker TensorFlow container, see https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/tensorflow/README.rst.

Use TensorFlow Legacy Mode

Use legacy mode TensorFlow training scripts to run TensorFlow jobs in Amazon SageMaker if:

For iformation about writing legacy mode TensorFlow scipts to use with the Amazon SageMaker Python SDK, see https://github.com/aws/sagemaker-python-sdk/tree/v1.12.0/src/sagemaker/tensorflow#tensorflow-sagemaker-estimators-and-models.