MNIST training with PyTorch

MNIST is a widely used dataset for handwritten digit classification. It consists of 70,000 labeled 28x28 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits). This tutorial will show how to train and test an MNIST model on SageMaker using PyTorch.

[ ]:
import os
import json

import sagemaker
from sagemaker.pytorch import PyTorch
from sagemaker import get_execution_role

sess = sagemaker.Session()

role = get_execution_role()

output_path = "s3://" + sess.default_bucket() + "/mnist"

PyTorch Estimator

The PyTorch class allows you to run your training script on SageMaker infrastracture in a containerized environment. In this notebook, we refer to this container as training container.

You need to configure it with the following parameters to set up the environment:

  • entry_point: A user defined python file to be used by the training container as the instructions for training. We further discuss this file in the next subsection.

  • role: An IAM role to make AWS service requests

  • instance_type: The type of SageMaker instance to run your training script. Set it to local if you want to run the training job on the SageMaker instance you are using to run this notebook

  • instance count: The number of instances you need to run your training job. Multiple instances are needed for distributed training.

  • output_path: S3 bucket URI to save training output (model artifacts and output files)

  • framework_version: The version of PyTorch you need to use.

  • py_version: The python version you need to use

For more information, see the API reference

Implement the entry point for training

The entry point for training is a python script that provides all the code for training a PyTorch model. It is used by the SageMaker PyTorch Estimator (PyTorch class above) as the entry point for running the training job.

Under the hood, SageMaker PyTorch Estimator creates a docker image with runtime environemnts specified by the parameters you used to initiated the estimator class and it injects the training script into the docker image to be used as the entry point to run the container.

In the rest of the notebook, we use training image to refer to the docker image specified by the PyTorch Estimator and training container to refer to the container that runs the training image.

This means your training script is very similar to a training script you might run outside Amazon SageMaker, but it can access the useful environment variables provided by the training image. Checkout the short list of environment variables provided by the SageMaker service to see some common environment variables you might used. Checkout the complete list of environment variables for a complete description of all environment variables your training script can access to.

In this example, we use the training script code/ as the entry point for our PyTorch Estimator.

[ ]:
!pygmentize 'code/'

Set hyperparameters

In addition, PyTorch estimator allows you to parse command line arguments to your training script via hyperparameters.

Note: local mode is not supported in SageMaker Studio

[ ]:
# set local_mode to be True if you want to run the training script
# on the machine that runs this notebook

local_mode = False

if local_mode:
    instance_type = "local"
    instance_type = "ml.c4.xlarge"

est = PyTorch(
    source_dir="code",  # directory of your training script
    hyperparameters={"batch-size": 128, "epochs": 1, "learning-rate": 1e-3, "log-interval": 100},

The training container executes your training script like

python --batch-size 100 --epochs 1 --learning-rate 1e-3 \
    --log-interval 100

Set up channels for training and testing data

You need to tell PyTorch estimator where to find your training and testing data. It can be a link to an S3 bucket or it can be a path in your local file system if you use local mode. In this example, we download the MNIST data from a public S3 bucket and upload it to your default bucket.

[ ]:
import logging
import boto3
from botocore.exceptions import ClientError

# Download training and testing data from a public S3 bucket

def download_from_s3(data_dir="/tmp/data", train=True):
    """Download MNIST dataset and convert it to numpy array

        data_dir (str): directory to save the data
        train (bool): download training set


    # Get global config
    with open("code/config.json", "r") as f:
        CONFIG = json.load(f)

    if not os.path.exists(data_dir):

    if train:
        images_file = "train-images-idx3-ubyte.gz"
        labels_file = "train-labels-idx1-ubyte.gz"
        images_file = "t10k-images-idx3-ubyte.gz"
        labels_file = "t10k-labels-idx1-ubyte.gz"

    # download objects
    s3 = boto3.client("s3")
    bucket = CONFIG["public_bucket"]
    for obj in [images_file, labels_file]:
        key = os.path.join("datasets/image/MNIST", obj)
        dest = os.path.join(data_dir, obj)
        if not os.path.exists(dest):
            s3.download_file(bucket, key, dest)

download_from_s3("/tmp/data", True)
download_from_s3("/tmp/data", False)
[ ]:
# upload to the default bucket

prefix = "mnist"
bucket = sess.default_bucket()
loc = sess.upload_data(path="/tmp/data", bucket=bucket, key_prefix=prefix)

channels = {"training": loc, "testing": loc}

The keys of the dictionary channels are parsed to the training image and it creates the environment variable SM_CHANNEL_<key name>.

In this example, SM_CHANNEL_TRAINING and SM_CHANNEL_TESTING are created in the training image (checkout how code/ access these variables). For more information, see: SM_CHANNEL_{channel_name}

If you want, you can create a channel for validation:

channels = {
    'training': train_data_loc,
    'validation': val_data_loc,
    'test': test_data_loc

You can then access this channel within your training script via SM_CHANNEL_VALIDATION

Run the training script on SageMaker

Now, the training container has everything to execute your training script. You can start the container by calling fit method.

[ ]:

Inspect and store model data

Now, the training is finished, the model artifact has been saved in the output_path. We

[ ]:
pt_mnist_model_data = est.model_data
print("Model artifact saved at:\n", pt_mnist_model_data)

We store the variable model_data in the current notebook kernel. In the next notebook, you will learn how to retrieve the model artifact and deploy to a SageMaker endpoint.

[ ]:
%store pt_mnist_model_data

Test and debug the entry point before executing the training container

The entry point code/ provided here has been tested and it can be executed in the training container. When you do develop your own training script, it is a good practice to simulate the container environment in the local shell and test it before sending it to SageMaker, because debugging in a containerized environment is rather cumbersome. The following script shows how you can test your training script:

[ ]:
!pygmentize code/