SageMaker HPO with MLflow

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.

Train a PyTorch model using HPO in SageMaker and track with MLflow using nested runs

Setup environment

Install necessary libraries

Import necessary libraries

Declare some variables used later

Get some training data

Download MNIST data

Upload data to S3

Write your training script

Since we're using MLflow in our training script, let's make sure the container installs mlflow along with our MLflow AWS plugin before running our training script. We can do this by creating a requirements.txt file and putting it in the same directory as our training script.

SageMaker HPO and MLflow

Create a MLflow experiment called MNIST. We'll give this SageMaker HPO job a run name, HPODemo. Each training attempt will be its own child run under HPODemo.

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.