Track an experiment while training a Pytorch model with a SageMaker Training Job

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

This notebook shows how you can use the SageMaker SDK to track a Machine Learning experiment using a Pytorch model trained in a SageMaker Training Job with Script mode, where you will provide the model script file.

We introduce two concepts in this notebook -

• Experiment: An experiment is a collection of runs. When you initialize a run in your training loop, you include the name of the experiment that the run belongs to. Experiment names must be unique within your AWS account.
• Run: A run consists of all the inputs, parameters, configurations, and results for one iteration of model training. Initialize an experiment run for tracking a training job with Run().

To execute this notebook in SageMaker Studio, you should select the PyTorch 1.12 Python 3.8 CPU Optimizer image.

You can track artifacts for experiments, including datasets, algorithms, hyperparameters and metrics. Experiments executed on SageMaker such as SageMaker training jobs are automatically tracked and any existen SageMaker experiment on your AWS account is automatically migrated to the new UI version.

In this notebook we will demonstrate the capabilities through an MNIST handwritten digits classification example. The notebook is organized as follow:

1. Train a Convolutional Neural Network (CNN) Model and log the model training metrics
2. Tune the hyperparameters that configures the number of hidden channels and the optimized in the model. Track teh parameter's configuration, resulting model loss and accuracy and automatically plot a confusion matrix using the Experiments capabilities of the SageMaker SDK.
3. Analyse your model results and plot graphs comparing your model different runs generated from the tunning step 3.

Runtime

This notebook takes approximately 45 minutes to run.

Contents

1. Install modules
2. Setup
3. Create model training script
4. Train model with Run context
5. Contact

Install modules

Let's ensure we have the latest SageMaker SDK available, including the SageMaker Experiments functionality

Setup

Import required libraries and set logging and experiment configuration

SageMaker Experiments now provides the Run class that allows you to create a new experiment run.

Create model training script

Let's create mnist.py, the pytorch script file to train our model.

The cell above saves the mnist.py file to our script folder. The file implements the code necessary to train our PyTorch model in SageMaker, using the SageMaker PyTorch image. It uses the load_run function to automatically detect the experiment configuration and run.log_parameter, run.log_parameters, run.log_file, run.log_metric and run.log_confusion_matrix to track the model training

Train model with Run context

Let's now train the model with passing the experiement run context to the training job

Checking the SageMaker Experiments UI, you can observe the Experiment run, populated with the metrics and parameters logged. We can also see the automatically generated outputs for the model data

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.