Compile and Train a Vision Transformer Model on the Caltech-256 Dataset using a Single Node

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.

1. Introduction
2. Development Environment and Permissions
   1. Installation
   2. SageMaker environment
3. Working with the Caltech-256 dataset
4. SageMaker Training Job
   1. Training Setup
   2. Training with Native TensorFlow
   3. Training with Optimized TensorFlow
5. Analysis
   1. Savings from Training Compiler
   2. Convergence of Training
6. Clean up

SageMaker Training Compiler Overview

SageMaker Training Compiler is a capability of SageMaker that makes hard-to-implement optimizations to reduce training time on GPU instances. The compiler optimizes Deep Learning (DL) models to accelerate training by more efficiently using SageMaker machine learning (ML) GPU instances. SageMaker Training Compiler is available at no additional charge within SageMaker and can help reduce total billable time as it accelerates training.

SageMaker Training Compiler is integrated into the AWS Deep Learning Containers (DLCs). Using the SageMaker Training Compiler enabled AWS DLCs, you can compile and optimize training jobs on GPU instances with minimal changes to your code. Bring your deep learning models to SageMaker and enable SageMaker Training Compiler to accelerate the speed of your training job on SageMaker ML instances for accelerated computing.

For more information, see SageMaker Training Compiler in the Amazon SageMaker Developer Guide.

Introduction

In this demo, you'll use Amazon SageMaker Training Compiler to train the Vision Transformer model on the Caltech-256 dataset. To get started, we need to set up the environment with a few prerequisite steps, for permissions, configurations, and so on.

NOTE: You can run this demo in SageMaker Studio, SageMaker notebook instances, or your local machine with AWS CLI set up. If using SageMaker Studio or SageMaker notebook instances, make sure you choose one of the TensorFlow-based kernels, Python 3 (TensorFlow x.y Python 3.x CPU Optimized) or conda_tensorflow_p39 respectively.

NOTE: This notebook uses a ml.p3.2xlarge instance with a single GPU. However, it can easily be extended to multiple GPUs on a single node. If you don't have enough quota, see Request a service quota increase for SageMaker resources.

Development Environment

Installation

This example notebook requires SageMaker Python SDK v2.115.0 or later

SageMaker environment

Working with the Caltech-256 dataset

We have hosted the Caltech-256 dataset in S3 in us-east-1. We will transfer this dataset to your account and region for use with SageMaker Training.

The dataset consists of JPEG images organized into directories with each directory representing an object category.

SageMaker Training Job

To create a SageMaker training job, we use a TensorFlow estimator. Using the estimator, you can define which training script should SageMaker use through entry_point, which instance_type to use for training, which hyperparameters to pass, and so on.

When a SageMaker training job starts, SageMaker takes care of starting and managing all the required machine learning instances, picks up the TensorFlow Deep Learning Container, uploads your training script, and downloads the data from sagemaker_session_bucket into the container at /opt/ml/input/data.

In the following section, you learn how to set up two versions of the SageMaker TensorFlow estimator, a native one without the compiler and an optimized one with the compiler.

Training with Native TensorFlow

The BATCH_SIZE in the following code cell is the maximum batch that can fit into the memory of a ml.p3.2xlarge instance while giving the best training speed. If you change the model, instance type, and other parameters, you need to do some experiments to find the largest batch size that will fit into GPU memory.

Set EPOCHS to the number of times you would like to loop over the training data.

Training with Optimized TensorFlow

Compilation through Training Compiler changes the memory footprint of the model. Most commonly, this manifests as a reduction in memory utilization and a consequent increase in the largest batch size that can fit on the GPU. But in some cases the compiler intelligently promotes caching which leads to a decrease in the largest batch size that can fit on the GPU. Note that if you want to change the batch size, you must adjust the learning rate appropriately.

Note: We recommend you to turn the SageMaker Debugger's profiling and debugging tools off when you use compilation to avoid additional overheads.

Wait for training jobs to complete

The training jobs described above typically take around 40 mins to complete

Note: If the estimator object is no longer available due to a kernel break or refresh, you need to directly use the training job name and manually attach the training job to a new TensorFlow estimator. For example:

native_estimator = TensorFlow.attach("<your_training_job_name>")

Analysis

Here we view the training metrics from the training jobs as a Pandas dataframe

Native TensorFlow

Optimized TensorFlow

Savings from Training Compiler

Let us calculate the actual savings on the training jobs above and the potential for savings for a longer training job.

Actual Savings

To get the actual savings, we use the describe_training_job API to get the billable seconds for each training job.

Potential savings

The Training Compiler works by compiling the model graph once per input shape and reusing the cached graph for subsequent steps. As a result the first few steps of training incur an increased latency owing to compilation which we refer to as the compilation overhead. This overhead is amortized over time thanks to the subsequent steps being much faster. We will demonstrate this below.

We calculate the potential savings below from the difference in steady state epoch latency between native TensorFlow and optimized TensorFlow

Convergence of Training

Training Compiler brings down total training time by intelligently choosing between memory utilization and core utilization in the GPU. This does not have any effect on the model arithmetic and consequently convergence of the model.

However, since we are working with a new batch size, hyperparameters like - learning rate, learning rate schedule and weight decay might have to be scaled and tuned for the new batch size

We can see that the model's convergence behavior is similar with and without Training Compiler. Here we have tuned the batch size specific hyperparameters - Learning Rate and Weight Decay using a linear scaling.

Learning rate is directly proportional to the batch size:

new_learning_rate = old_learning_rate * new_batch_size/old_batch_size

Weight decay is inversely proportional to the batch size:

new_weight_decay = old_weight_decay * old_batch_size/new_batch_size

Better results can be achieved with further tuning. Check out Automatic Model Tuning for tuning.

Clean up

Stop all training jobs launched if the jobs are still running.

Also, to find instructions on cleaning up resources, see Clean Up in the Amazon SageMaker Developer Guide.

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.