Amazon SageMaker Automatic Model Tuning becomes more efficient with warm start of hyperparameter tuning jobs

Earlier this year, we launched Amazon SageMaker Automatic Model Tuning, which allows developers and data scientists to save significant time and effort in training and tuning their machine learning models. Today, we are launching warm start of hyperparameter tuning jobs in Automatic Model Tuning. Data scientists and developers can now create a new hyperparameter tuning job based on selected parent jobs, so that training jobs conducted in those parent jobs can be reused as prior knowledge. Warm start of hyperparameter tuning jobs will accelerate the hyperparameter tuning process and reduce the cost for tuning models.

While data scientists and developers could already efficiently tune their models through Automatic Model Tuning, there are still places where they need more help. For example, they might start a hyperparameter tuning job with a small budget, and, after analyzing the results, decide that they want to continue tuning the model with a larger budget. Potentially they might use different hyperparameter configurations (e.g., by adding more hyperparameters to tune or trying different search ranges for some hyperparameters). Another example is when data scientists or developers might want to re-tune a model after they have collected new data subsequent to a previous model tuning. In both cases, starting a hyperparameter tuning job with prior knowledge collected from previous tuning jobs on this model can help get to the best model faster, and end up saving cost for customers. However, previously every tuning job would start from scratch. Even if the same model was already tuned with a similar tuning configuration, no information was reused.

Warm start of hyperparameter tuning jobs addresses these needs. Now we’ll show you how to iteratively tune your model leveraging warm start.

Tuning an image classification model leveraging warm start

In this example, we’ll build an image classifier and iteratively tune it by running multiple hyperparameter tuning jobs leveraging warm start. We’ll use the Amazon SageMaker built-in image classification algorithm and train the model against the Caltech-256 dataset. You can find the full sample notebook here.

Set up and launch the hyperparameter tuning job

We’ll skip the steps of creating a notebook instance, preparing the dataset, and pushing it to Amazon S3, and directly start from launching a hyperparameter tuning job. The sample notebook has all the details so we won’t go through the process here.

We’ll run this first tuning job to learn about the search space and evaluate the impact of tuning tunable hyperparameters in image classification. This job will assess if tuning the model is promising, and if we want to continue the tuning by creating a subsequent tuning job.

To create a tuning job, we first need to create a training estimator for the built-in image classification algorithm, and specify values for every hyperparameter of this algorithm, except for those we plan to tune. To learn more about hyperparameters of the built-in image classification algorithm, you can explore our documentation.

s3_output_location = 's3://{}/{}/output'.format(bucket, prefix)
s3_input_train = sagemaker.s3_input(s3_data='s3://{}/{}/train'.format(bucket, prefix), content_type='application/x-recordio')
s3_input_validation = sagemaker.s3_input(s3_data='s3://{}/{}/validation/'.format(bucket, prefix), content_type='application/x-recordio')
sess = sagemaker.Session()

imageclassification = sagemaker.estimator.Estimator(training_image,
 role, 
 train_instance_count=1, 
 train_instance_type='ml.p3.8xlarge',
 output_path=s3_output_location,
 sagemaker_session=sess)

imageclassification.set_hyperparameters(num_layers=18,
 image_shape='3,224,224',
 num_classes=257,
 num_training_samples=15420,
 mini_batch_size=128,
 epochs=50,
 optimizer='sgd',
 top_k='2',
 precision_dtype='float32',
 augmentation_type='crop')

Now that we have the estimator, we can create a hyperparameter tuning job with the estimator and specify the search ranges for hyperparameters we want to tune and the number of total training jobs we want to run.

We selected the three hyperparameters that we believe are most likely to affect the model quality, and thus our objective metric. Since we don’t know yet the values that lead to the best model, we chose the full range of search for momentum and weight_decay as specified in image classification documentation, and a smaller range of search for learning rate (0.0001, 0.05):

learning_rate: controls how fast the training algorithm will try to optimize your model. Lower learning rates can achieve better accuracy but will take more time to train your model. Higher learning rates can fail to improve your model accuracy. You need to find a good balance for this attribute.
momentum: uses information from the direction of our previous update to inform our current update. The default value of 0 means weight updates are based only on the information in the current batch.
weight_decay: penalizes weights when they grow too large. The

default value of 0 means no penalty.

In this case we don’t need to specify the regular expressions for the objective metric because we are using one of the Amazon SageMaker built-in algorithms.

from sagemaker.tuner import IntegerParameter, CategoricalParameter, ContinuousParameter, HyperparameterTuner

hyperparameter_ranges = {'learning_rate': ContinuousParameter(0.0001, 0.05),
 'momentum': ContinuousParameter(0.0, 0.99),
 'weight_decay': ContinuousParameter(0.0, 0.99)}

objective_metric_name = 'validation:accuracy'

tuner = HyperparameterTuner(imageclassification,
 objective_metric_name,
 hyperparameter_ranges,
 objective_type='Maximize',
 max_jobs=10,
 max_parallel_jobs=2) 

After the hyperparameter tuning job finishes, we can bring in a table of metrics using the HyperparameterTuningJobAnalytics API action from the Amazon SageMaker Python SDK.

tuner_parent = sagemaker.HyperparameterTuningJobAnalytics(tuning_job_name)
tuner_parent.dataframe().sort_values(['FinalObjectiveValue'], ascending=False)

This table shows a subset of the training jobs that have been run. You can look at all of the results by running the notebook. Observe that the hyperparameters we are tuning have a significant impact on the objective metric values for the image classification algorithm. Choosing different values gives very different results.

Using the HPO_Analyze_TuningJob_Results.ipynb notebook, we can plot how the objective metric changes over time as the tuning job progresses.

You can see that the objective metric values improve over time as Automatic Model Tuning is learning through the search space. We might get further improvement beyond the 0.33 validation accuracy by running a few more training jobs. To validate the hypothesis, we‘ll run a second tuning job with another 10 training jobs. This time we‘ll use warm start to reuse the learning we gathered from the first tuning job.

Don‘t worry if youdon‘t get a trend as clear as the one we just discussed, given the nature of randomness in a tuning process. Even running the same experiment won‘t give you the same result, but typically you should see an overall trend of model quality improvement.

Set up and launch a hyperparameter tuning job using a warm start configuration

To use warm start in the new tuning job, we need to specify two parameters:

The list of parent tuning jobs the new tuning job should use as a starting point. (The maximum number of parents can be 5 but we will use 1 in this example.)

The type of warm start configuration:

IDENTICAL_DATA_AND_ALGORITHM warm starts a tuning job with previous evaluations essentially with the same task, allowing for slight changes in the search space. This option should be used when the data set and the algorithm container haven’t changed.
TRANSFER_LEARNING warm starts a tuning job with the evaluations from similar tasks, allowing both search space, algorithm image, and dataset change.

In this example we’ll use IDENTICAL_DATA_AND_ALGORITHM because we are not changing the data set or algorithm, we are just running more training jobs.

We will use the Amazon SageMaker console to launch our second tuning job with warm start. Open the Amazon SageMaker console, and in the left navigation pane choose Training.-Then choose Hyperparameter tuning jobs and Create hyperparameter tuning job. At the top of the page, enable Warm start with identical data and algorithm Warm start type. The next step is to select the parent jobs of the new tuning job:

The console allows us to easily populate the values of the new tuning job by using Copy settings from the parent tuning job. After choosing Copy settings, the form gets populated. Choose Next and validate that the static and tunable hyperparameters look good:

In this case, we are not changing any hyperparameter values, so we just need to choose Next again and create the new tuning job using warm start. Really simple!

After the warm start hyperparameter tuning job has completed, we can go back to the notebook to use tuner.analytics() to visualize how the objective metric changes over time for the parent tuning job (black data points) and the new tuning job we launched using warm start (red data points).

You can see that the new tuning job managed to find good hyperparameter configurations very early on, thanks to the prior knowledge from the parent tuning job. As the optimization continues, the objective metric continues improving and it reaches 0.47, which is significantly higher than the metric we had gotten (0.33) when we ran the first tuning job from scratch.

Lastly, to demonstrate how you could apply transfer learning to a tuning job using warm start, we’ll run a third tuning job using more data augmentations in the data set to see if those drive our validation accuracy further up. To apply more data augmentations we can use augmentation_type hyperparameter exposed by the Amazon SageMaker pre-built image classification algorithm. We’ll apply crop_color_transform transformation to the data set during training. With this transformation, in addition to crop and color transformations, random transformations (including rotation, shear, and aspect ratio variations) are applied to the image.

To create our last hyperparameter tuning job, we will use Transfer learning WarmStartType since our data set is going to change as a result of applying new data augmentations. We’ll use both of the two previous tuning jobs that we ran as parent tuning jobs and run 10 more training jobs. Let’s go back to the notebook to launch this last hyperparameter tuning job:

from sagemaker.tuner import WarmStartConfig, WarmStartTypes

parent_tuning_job_name_2 = warmstart_tuning_job_name
transfer_learning_config = WarmStartConfig(WarmStartTypes.TRANSFER_LEARNING, 
 parents={parent_tuning_job_name,parent_tuning_job_name_2})

imageclassification.set_hyperparameters(num_layers=18,
 image_shape='3,224,224',
 num_classes=257,
 num_training_samples=15420,
 mini_batch_size=128,
 epochs=50,
 optimizer='sgd',
 top_k='2',
 precision_dtype='float32',
 augmentation_type='crop_color_transform')

tuner_transfer_learning = HyperparameterTuner(imageclassification,
 objective_metric_name,
 hyperparameter_ranges,
 objective_type='Maximize',
 max_jobs=10,
 max_parallel_jobs=2,
 base_tuning_job_name='transferlearning',
 warm_start_config=transfer_learning_config)

tuner_transfer_learning.fit({'train': s3_input_train, 'validation': s3_input_validation},include_cls_metadata=False)

One last time, after the new hyperparameter tuning job has been completed, we can go use tuner.analytics() to visualize how the objective metric changed over time for the parent tuning jobs (black and red data points) and the new tuning job we launched using warm start transfer learning (blue data points).

After the tuning job has been completed, the objective metric has improved again and has reached 0.52.

If you are satisfied with the results, you can find the training job that generated the best model by getting BestTrainingJob in the Automatic Model Tuning describe API or by going to the console. From the console you can deploy the model to an Amazon SageMaker hosting endpoint.

Conclusion

To recap, we explored one use case that showed how using warm start can help explore the search space iteratively without losing the learning gathered in previous iterations. We also demonstrated how you can use warm start to transfer the learning of previous tuning jobs even if your dataset or algorithm has been changed, but you believe they are close enough to datasets or algorithms used in previous hyperparameter tuning jobs.

Warm start of hyperparameter tuning jobs is now available in all the AWS Regions where Amazon SageMaker is available today. For more information on Amazon SageMaker Automatic Model Tuning, visit Amazon SageMaker documentation.

About the Authors

Patricia Grao is a Software Development Manager in Amazon AI. She became passionate about machine learning while working in search ranking and query understanding in Amazon Search. She was part of the team that launched Amazon SageMaker Automatic Model Tuning.

Fela Winkelmolen works as an applied scientists for Amazon AI and was part of the team that launched the Automatic Model Tuning feature of Amazon SageMaker

Fan Li is a Product Manager of Amazon SageMaker. He used to be a big fan of ballroom dance but now loves whatever his 8-year-old son likes.