SageMaker and TensorFlow 2.0

EDIT: Amazon SageMaker does now support TF 2.0 and higher.

  • SageMaker + TensorFlow docs: https://sagemaker.readthedocs.io/en/stable/frameworks/tensorflow/using_tf.html
  • Supported Tensorflow versions (and Docker URIs): https://aws.amazon.com/releasenotes/available-deep-learning-containers-images

Original answer

Here is an example Dockerfile that uses the underlying SageMaker Containers library (this is what is used in the official pre-built Docker images):

FROM tensorflow/tensorflow:2.0.0b1

RUN pip install sagemaker-containers

# Copies the training code inside the container
COPY train.py /opt/ml/code/train.py

# Defines train.py as script entrypoint
ENV SAGEMAKER_PROGRAM train.py

For more information on this approach, see https://docs.aws.amazon.com/sagemaker/latest/dg/build-container-to-train-script-get-started.html


Update on 10 nov 2019:

There is now a way to use Tensorflow 2 in SageMaker, event though there is no shortcut to start TF 2 directly from SageMaker console.

  1. Start a conda Python3 Kernel

  2. Make some updates (one in each code cell):

!pip install --upgrade pip         # pip 19.0 or higher is required for TF 2
!pip install --upgrade setuptools  # Otherwise you'll get annoying warnings about bad installs
  1. Install Tensorflow 2
!pip install --user --upgrade tensorflow

Given the doc, this will install in $HOME.

Nota:

If you are using a GPU based instance of SageMaker, replace tensorflow by tensorflow-gpu.

You now can use TF 2 in your instance. This only needs to be done once, as long as the instance remains up.

To test, just run in the next cell:

import tensorflow as tf
print(tf.__version__)

You should see 2.0.0 or higher.