Sign In to the Console
Amazon SageMaker
Developer Guide

AWS Documentation » Amazon SageMaker » Developer Guide » Build Models » Choose an Algorithm » Use Your Own Algorithms or Models with Amazon SageMaker » Prebuilt Amazon SageMaker Docker Images for TensorFlow, MXNet, Chainer, and PyTorch

Prebuilt Amazon SageMaker Docker Images for TensorFlow, MXNet, Chainer, and PyTorch

Amazon SageMaker provides prebuilt Docker images that include deep learning framework libraries and other dependencies needed for training and inference. With the SageMaker Python SDK, you can train and deploy models using one of these popular deep learning frameworks. For instructions on installing and using the SDK, see Amazon SageMaker Python SDK.

The following table provides links to the GitHub repositories that contain the source code and Dockerfiles for each framework and for TensorFlow and MXNet Serving. The instructions linked are for using the Python SDK estimators to run your own training algorithms on Amazon SageMaker and your own models on Amazon SageMaker hosting.

Framework Prebuilt Docker Image Source Code Instructions
TensorFlow

Amazon SageMaker TensorFlow Containers

Amazon SageMaker TensorFlow Serving Container

Using TensorFlow with the SageMaker Python SDK
MXNet

Amazon SageMaker MXNet Containers

Amazon SageMaker MXNet Serving Container

Using MXNet with the SageMaker Python SDK
Chainer

Amazon SageMaker Chainer SageMaker Containers

Chainer SageMaker Estimators and Models
PyTorch

Amazon SageMaker PyTorch Containers

SageMaker PyTorch Estimators and Models

If you are not using the Amazon SageMaker Python SDK and one of its estimators to manage the container, you have to retrieve the relevant pre-built container. The Amazon SageMaker prebuilt Docker images are stored in Amazon Elastic Container Registry (Amazon ECR). To pull an image from an Amazon ECR repo or to push an image to an Amazon ECR repo, use the fullname registry address of the image. Amazon SageMaker uses the following URL patterns for the container image registery addresses:

<account_id>.dkr.ecr.<region>.amazonaws.com/<ECR repo name>:<framework version>-<processing unit type>-<python version>

Use the following commands to pull these images:

$(aws ecr get-login --no-include-email --registry-ids <account_id>)
docker pull <account_id>.dkr.ecr.<region>.amazonaws.com/<ECR repo name>:<framework version>-<processing unit type>-<python version>

The following table itemizes the supported values for each of the components in the URL registry addresses and how they are associated.

Here, for example, are some of the most common use cases for the deep learning frameworks supported the Amazon SageMaker:

  • If you want to use TensorFlow 1.13 or later to train a model:

    763104351884.dkr.ecr.<region>.amazonaws.com/tensorflow-training:1.13-gpu

  • If you want to use TensorFlow 1.14 or later to train a model:

    763104351884.dkr.ecr.<region>.amazonaws.com/tensorflow-training:1.14-gpu

  • If you want to use TensorFlow 1.14 or later for inference:

    763104351884.dkr.ecr.<region>.amazonaws.com/tensorflow-inference:1.14-gpu

  • If you want to use MxNet 1.4.1 or later to train a model or for inference with Python 3:

    763104351884.dkr.ecr.<region>.amazonaws.com/mxnet-training:1.4.1-gpu-py3

    763104351884.dkr.ecr.<region>.amazonaws.com/mxnet-inference:1.4.1-gpu-py3

  • If you want to use PyTorch 1.2.0 or later to train a model or for inference with Python 3:

    763104351884.dkr.ecr.<region>.amazonaws.com/pytorch-training:1.2.0-gpu-py3

    763104351884.dkr.ecr.<region>.amazonaws.com/pytorch-inference:1.2.0-gpu-py3

  • If you want to use Chainer to train a model or for inference with Python 2 or 3:

    520713654638.dkr.ecr.<region>.amazonaws.com/sagemaker-chainer:5.0.0-gpu-<python version>

  • If you want to use EIA containers for inference with TensorFlow 1.14 with Python 3 or for MxNet 1.4.1 with Python 2 or 3:

    763104351884.dkr.ecr.<region>.amazonaws.com/tensorflow-inference-eia:1.14-cpu

    763104351884.dkr.ecr.<region>.amazonaws.com/mxnet-inference-eia:1.4.1-cpu-<python version>

URL Component Description Supported Values
<account_id>

Specifies the ID for Amazon SageMaker accounts that contain the pre-built containers.

  • 763104351884

  • 520713654638

  • 871362719292

  • 057415533634

<region>

Specifies the AWS regions that contain the Amazon SageMaker accounts.

  • Accounts 763104351884 and 520713654638 are located in:

    us-west-1, us-west-2, us-east-1, us-east-2, ap-northeast-1, ap-northeast-2, ap-southeast-1, ap-southeast-2, ap-south-1, eu-west-1, eu-west-2, eu-central-1, ca-central-1

  • Accounts 871362719292 and 057415533634 are located in:

    ap-east-1

  • Account 763104351884 has EIA containers located in:

    us-west-2, us-east-1, us-east-2, eu-west-1, ap-northeast-1, ap-northeast-2

<ECR repo name>

Specifies the name of the public repository owned by Amazon SageMaker in the Amazon ECR.

Python 3 containers for TensorFlow-1.13 and later, for MXNet-1.4.1, and for PyTorch 1.2.0 and later in accounts 763104351884 and 871362719292:

  • tensorflow-training (also for Python 2)

  • tensorflow-inference

  • mxnet-training

  • mxnet-inference

  • pytorch-training (also for Python 2)

  • pytorch-inference (also for Python 2)

Other Python 2 and Python 3 containers in accounts 520713654638 and 057415533634:

  • sagemaker-tensorflow-scriptmode (for Python 2 only)

  • sagemaker-tensorflow-serving-eia

  • sagemaker-mxnet (for Python 2 only)

  • sagemaker-mxnet-serving (for Python 2 only)

  • sagemaker-mxnet-serving-eia

  • sagemaker-chainer

  • sagemaker-pytorch

Python 3 EIA containers for TensorFlow-1.14 and later and Python 2 and Python 3 EIA containers for MXNet-1.4.1 and later that use elastic inference are in account 763104351884:

  • tensorflow-inference-eia

  • mxnet-inference-eia
<framework version>

Specifies the framework and links to documentation for the estimators for each of the frameworks that explains how to specify the supported versions.

<processing unit type>

Specifies whether to use a GPU or CPU for training or hosting.

  • cpu (EIA images also use cpu in the label.)

  • gpu
<python version>

Specifies the version of Python used. (This tag is not used for the tensorflow-inference, tensorflow-inference-eia, and tensorflow-serving containers.)

  • py2

  • py3

Amazon SageMaker also provides prebuilt Docker images for scikit-learn and Spark ML. For information about Docker images that enable using scikit-learn and Spark ML solutions in Amazon SageMaker, see Prebuilt Amazon SageMaker Docker Images for Scikit-learn and Spark ML .

You can use prebuilt containers to deploy your custom models or models that you have purchased on AWS Marketplace that have been trained in a framework other than Amazon SageMaker. For an overview of the process of bringing the trained model artifacts into Amazon SageMaker and hosting them at an endpoint, see Bring Your Own Pretrained MXNet or TensorFlow Models into Amazon SageMaker.

You can customize these prebuilt containers or extend them to handle any additional functional requirements for your algorithm or model that the prebuilt Amazon SageMaker Docker image doesn't support. For an example, see Extending Our PyTorch Containers.