Transfer learning for TensorFlow image classification models in Amazon SageMaker

Amazon SageMaker provides a set of built-in algorithms, pre-trained models, and pre-built solution templates to help data scientists and machine learning (ML) practitioners get started quickly with training and deploying ML models. These algorithms and models can be used for both supervised and unsupervised learning. It can handle different types of input data such as tabular, image, and text.

Starting today, SageMaker offers a new built-in algorithm for image classification: Image Classification – TensorFlow. A supervised learning algorithm that supports transfer learning of many pretrained models available in . TensorFlow Hub. It takes an image as input and outputs the probability for each class label. Transfer learning can be used to fine-tune these pre-trained models even when large amounts of training images are not available. This is available via SageMaker built-in algorithms and the SageMaker JumpStart UI within Amazon SageMaker Studio. For more information, see the documentation Image Classification – TensorFlow and sample notebooks. Introducing SageMaker TensorFlow – Image Classification.

SageMaker’s Image Classification with TensorFlow provides transfer learning with many pre-trained models available on TensorFlow Hub. A classification layer is added to the pretrained TensorFlow Hub model depending on the number of class labels in the training data. The classification layer consists of a dropout layer and a dense layer, which is a fully connected layer with a 2-norm regularizer initialized with random weights. The model training has hyperparameters for the dropout rate of the dropout layer and the L2 regularization factor of the dense layer. The entire network, including pretrained models, or just the top classification layer can then be fine-tuned with new training data. With this transfer learning mode, training can be achieved even on smaller datasets.

How to use the new TensorFlow image classification algorithm

In this section, we will use the TensorFlow image classification algorithm SageMaker Python SDKSee SageMaker JumpStart for how to use it from the Studio UI.

This algorithm supports transfer learning of pretrained models listed in TensorFlow Hub Models.Each model is unique model_idThe following code shows how to tweak MobileNet V2 1.00 224 identified by . model_id tensorflow-ic-imagenet-mobilenet-v2-100-224-classification-4 On a custom training dataset.about each model_idto launch the SageMaker training job, estimator To use the classes from the SageMaker Python SDK, you need to use utility functions provided by SageMaker to get the Docker image URI, training script URI, and pretrained model URI. The training script URI contains all the code required for data processing, loading pretrained models, model training, and saving trained models for inference. The pretrained model URI contains the pretrained model architecture definition and model parameters. Note that the Docker image URI and training script URI are the same for all TensorFlow image classification models. A pretrained model URI is specific to a particular model. The trained model tarball is pre-downloaded from TensorFlow Hub and stored in an Amazon Simple Storage Service (Amazon S3) bucket with the proper model signature, so the training job runs with network isolation. See the code below.

from sagemaker import image_uris, model_uris, script_uris
from sagemaker.estimator import Estimator

model_id, model_version = "tensorflow-ic-imagenet-mobilenet-v2-100-224-classification-4", "*"
training_instance_type = "ml.p3.2xlarge"

# Retrieve the docker image
train_image_uri = image_uris.retrieve(model_id=model_id,model_version=model_version,image_scope="training",instance_type=training_instance_type,region=None,framework=None)

# Retrieve the training script
train_source_uri = script_uris.retrieve(model_id=model_id, model_version=model_version, script_scope="training")

# Retrieve the pre-trained model tarball for transfer learning
train_model_uri = model_uris.retrieve(model_id=model_id, model_version=model_version, model_scope="training")

output_bucket = sess.default_bucket()
output_prefix = "jumpstart-example-ic-training"
s3_output_location = f"s3://{output_bucket}/{output_prefix}/output"

Using these model-specific training artifacts, estimator class:

# Create SageMaker Estimator instance
tf_ic_estimator = Estimator(
    role=aws_role,
    image_uri=train_image_uri,
    source_dir=train_source_uri,
    model_uri=train_model_uri,
    entry_point="transfer_learning.py",
    instance_count=1,
    instance_type=training_instance_type,
    max_run=360000,
    hyperparameters=hyperparameters,
    output_path=s3_output_location,
)

Second, for transfer learning on custom datasets, you may need to change the default values ​​of the training hyperparameters. HyperparameterTo get a Python dictionary of these hyperparameters with their default values, call: hyperparameters.retrieve_defaultUpdate it as needed and pass it to the Estimator class. Note that the default values ​​of some hyperparameters are model dependent. For large models, the default batch size is small and train_only_top_layer Hyperparameters are set True. Hyperparameters Train_only_top_layer Defines which model parameters are changed during the fine-tuning process.If train_only_top_layer teeth True, the parameters of the classification layer are changed and the remaining parameters remain constant during the fine-tuning process. on the other hand, train_only_top_layer teeth False, with all parameters of the model fine-tuned. See the code below.

from sagemaker import hyperparameters
# Retrieve the default hyper-parameters for fine-tuning the model
hyperparameters = hyperparameters.retrieve_default(model_id=model_id, model_version=model_version)

# [Optional] Override default hyperparameters with custom values
hyperparameters["epochs"] = "5"

The following code provides a default training dataset hosted in an S3 bucket. we, tf_flowers dataset as the default dataset for fine-tuning the model. The dataset consists of images of five types of flowers.The dataset was downloaded from TensorFlow under Apache 2.0 License.

# Sample training data is available in this bucket
training_data_bucket = f"jumpstart-cache-prod-{aws_region}"
training_data_prefix = "training-datasets/tf_flowers/"

training_dataset_s3_path = f"s3://{training_data_bucket}/{training_data_prefix}"

Finally, to launch a SageMaker training job to fine-tune your model, call: .fit Pass the S3 location of the training dataset in an object of the Estimator class.

# Launch a SageMaker Training job by passing s3 path of the training data
tf_ic_estimator.fit({"training": training_dataset_s3_path}, logs=True)

For more information on how to use the new SageMaker TensorFlow image classification algorithm for transfer learning on custom datasets, see Deploying a fine-tuned model, running inference on the deployed model, and without first fine-tuning Deploy the pre-trained model as-is. See the example notebook below for a custom dataset. Introducing SageMaker TensorFlow – Image Classification.

Input/output interface for TensorFlow image classification algorithms

You can fine-tune each pretrained model listed in . TensorFlow hub model to any dataset containing images belonging to any number of classes. The goal is to minimize the prediction error of the input data. The model returned by finetuning can be further deployed for inference. Here are the steps on how to format the training data for input to the model:

• input – A directory with as many subdirectories as there are classes. Each subdirectory must have images belonging to that class in .jpg, .jpeg, or .png format.
• output – A fine-tuned model that can be deployed for inference or further trained using incremental training. Pre- and post-processing signatures are added to the fine-tuned model, which takes a raw .jpg image as input and returns class probabilities. A file that maps class indices to class labels is saved with the model.

If your training data contains images of two classes, your input directory will look like the following example: roses When dandelionThe .S3 path looks like this: s3://bucket_name/input_directory/. Note the end / is necessary.folder name and roses, dandelion, and the .jpg filename can be anything. The label mapping file, which is stored with the trained model in an S3 bucket, maps the folder names roses and tandelion to indices in the list of class probabilities output by the model. The mapping follows alphabetical order by folder name. In the following example, index 0 in the model output list corresponds. dandelionand index 1. roses.

input_directory
    |--roses
        |--abc.jpg
        |--def.jpg
    |--dandelion
        |--ghi.jpg
        |--jkl.jpg

Inference with TensorFlow image classification algorithms

The generated model can be hosted for inference and supports encoded .jpg, .jpeg, and .png image formats. application/x-image Content type. Input images are automatically resized. The output contains the probability value, the class labels for all classes, and the predicted label corresponding to the class index with the highest probability, encoded in JSON format. A TensorFlow image classification model processes one image per request and outputs only one line in JSON. Here’s an example response in JSON:

accept: application/json;verbose

 {"probabilities": [prob_0, prob_1, prob_2, ...],
  "labels":        [label_0, label_1, label_2, ...],
  "predicted_label": predicted_label}

If accept is set to application/jsonIf , the model outputs probabilities only. For more information on training and inference, see the sample notebook. Introducing SageMaker TensorFlow – Image Classification.

Use SageMaker built-in algorithms via the JumpStart UI

You can also use SageMaker TensorFlow image classification and other built-in algorithms with just a few clicks from the JumpStart UI. JumpStart is a feature of SageMaker that allows you to train and deploy built-in algorithms and pre-trained models from various ML frameworks and model hubs through a graphical interface. You can also combine ML models with various other AWS services to deploy full-fledged ML solutions that solve targeted use cases. Check out Using TensorFlow Hub and Hugging Face Models to Perform Text Classification with Amazon SageMaker JumpStart to learn how to use JumpStart to train an algorithm or pre-trained model in just a few clicks .

Conclusion

In this post, we announced the release of the SageMaker TensorFlow image classification built-in algorithm. Using this algorithm, we have provided sample code on how to do transfer learning on a custom dataset using a pre-trained model in TensorFlow Hub.For more information see the documentation and sample notes.

About the author

Dr. Ashish Khetan is a Senior Applied Scientist using Amazon SageMaker built-in algorithms, helping develop machine learning algorithms. He received his doctorate from the University of Illinois at Urbana-Champaign. He is an active researcher in machine learning and statistical inference and has published many papers at NeurIPS, ICML, ICLR, JMLR, ACL, and his EMNLP his conferences.

Dr. Vivek Madan I am an Applied Scientist on the Amazon SageMaker JumpStart team. He earned his doctorate from the University of Illinois at Urbana, his school in Champaign, and served as a postdoctoral fellow at the Georgia Institute of Technology. He is an active researcher in machine learning and algorithm design and has published papers at EMNLP, ICLR, COLT, FOCS, and his SODA conferences.

Joan Moura AI/ML Specialist Solutions Architect at Amazon Web Services. He primarily focuses on his NLP use his case, helping customers optimize training and deployment of deep learning models. He is also an active proponent of low-code ML solutions and ML-specific hardware.

Large Pen Matcha Senior AI/ML Specialist Solutions Architect at AWS. He works with education, government, and non-profit customers on machine learning and artificial intelligence-related projects, helping them build solutions using AWS. When he’s not supporting his customers, he likes to travel to new places.