Using Shutterstock's Image datasets to train a multi-label image classification model

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.

Introduction

This example of multi-label image classification trains the Amazon SageMaker 1P image classification algorithm. We will use the Amazon SageMaker image classification algorithm in transfer learning mode to fine-tune a pre-trained model (trained on ImageNet data) to learn to classify a new multi-label dataset. The pre-trained model will be fine-tuned using the Free Sample: Images & Metadata of “Whole Foods” Shoppers dataset from Shutterstock’s Image Datasets.

You can learn more about this implementation from Using Shutterstock's image datasets to train your computer vision models.

To get started, we need to set up the environment with a few prerequisite steps, for permissions, configurations, and so on.

Note: This notebook is a modification of the existing SageMaker example notebook for multi-label image classification.

Prerequisites

Step 1: Set permissions and environment variables

Here we set up the pointers and authentication to AWS services. There are three parts to this:

• The roles used to give learning and hosting access to your data. This will automatically be obtained from the role used to start the notebook
• The S3 bucket that you want to use for SageMaker training and model data
• The S3 bucket that is storing the images from the Shutterstock Image Dataset for Whole Foods Shoppers
• The SageMaker image classification docker image (which we will not need to change)

FILL IN YOUR S3 BUCKET NAME: For our images_bucket variable, replace <enter-images-bucket-name> with the name of the S3 bucket that contains your Shutterstock Image Dataset. If you stored your images under a prefix, please also update the images_bucket_prefix variable with this prefix.

Data Preparation

The Shutterstock Image Datasets are collections of curated images from Shutterstock’s library of more than 370 million images. You can subscribe to one of the preexisting collections, such as Food & Beverage, Clothing, or Hospitality, or you can work with the Shutterstock Data Exchange team to request a custom collection of images for your use case. Each image includes a descriptive title with up to 200 characters and an optimal 7-50 keywords.

For this example, we use the Free Sample: Images & Metadata of “Whole Foods” Shoppers dataset from Shutterstock’s Image Datasets to demonstrate how to train a multi-label image classification model using Shutterstock’s pre-labeled image assets. This dataset contains images of Whole Foods shoppers. Each image is tagged with 30-50 keywords describing what is seen in the image.

An image can contain objects of multiple keywords. Because the purpose of this notebook is to show how the Shutterstock dataset can be used to train for multi-label image classification, we will limit the keywords that we are training on for demonstration purposes. We have a limited set of 200 images in our free sample.

To start, we will train the model to learn to classify the first 5 keywords that appear in our dataset.

1. Adult
2. Beautiful
3. Buying
4. Car
5. Casual

The image classification algorithm can take two types of input formats. The first is a recordIO file format and the other is an LST file format. We will use the LST file format for training.

Step 2: Parse the metadata .csv file to create a .lst file for training

A .lst file is a tab-separated file with three columns that contains a list of image files. The first column specifies the image index, the second column specifies the class label index for the image, and the third column specifies the relative path of the image file. The image index in the first column should be unique across all the images.

The Whole Foods Shoppers dataset comes with a metadata file in CSV format. We will use this to generate our .lst file. This CSV metadata file has the following headers in the first row: 'ASSET_ID' | 'COLLECTION' |'KEYWORDS' | 'DESCRIPTION'.

Step 3: Create the training and validation datasets

Create training and validation set by splitting the LST file. Limit images to the top 150 images. Use 70% of the images for training and 30% of the images for validation.

Step 4: Copy the training data to your SageMaker S3 Bucket

Upload the data into the SageMaker S3 bucket. The images are uploaded under the train and validation prefixes. The LST files are uploaded under train_lst and validation_lst prefixes.

Training our Multi-Label Image Classification Model

Now that we have prepared our training data, we are ready to train our image classification model. First, we will create a SageMaker Estimator object. This estimator will launch the training job.

Step 5: Set the training parameters

There are two kinds of parameters that need to be set for training. The first are the parameters for the training job. These include:

• Training instance count: This is the number of instances on which to run the training. When the number of instances is greater than one, then the image classification algorithm will run in distributed settings.
• Training instance type: This indicates the type of machine on which to run the training. Here, we will use ml.p2.xlarge.
• Output path: This the S3 folder in which the training output is stored

Step 6: Set the algorithm parameters

Apart from the above set of parameters, there are hyperparameters that are specific to the algorithm. These are:

• num_layers: The number of layers (depth) for the network.
• use_pretrained_model: Set to 1 to use pretrained model for transfer learning.
• image_shape: The input image dimensions,'num_channels, height, width', for the network. It should be no larger than the actual image size. The number of channels should be same as the actual image.
• num_classes: This is the number of output classes for the dataset.
• mini_batch_size: The number of training samples used for each mini batch. In distributed training, the number of training samples used per batch will be N * mini_batch_size where N is the number of hosts on which training is run
• resize: Resize the image before using it for training. The images are resized so that the shortest side is of this parameter. If the parameter is not set, then the training data is used as such without resizing.
• epochs: Number of training epochs
• learning_rate: Learning rate for training
• num_training_samples: This is the total number of training samples.
• use_weighted_loss: This parameter is used to balance the influence of the positive and negative samples within the dataset.
• augmentation_type: This parameter determines the type of augmentation used for training. It can take on three values, 'crop', 'crop_color' and 'crop_color_transform'
• precision_dtype: The data type precision used during training. Using float16 can lead to faster training with minimal drop in accuracy, particularly on P3 machines. By default, the parameter is set to float32
• multi_label: Set multi_label to 1 for multi-label processing

You can find a detailed description of all the algorithm parameters at https://docs.aws.amazon.com/sagemaker/latest/dg/IC-Hyperparameter.html

Step 7: Specify the input data

Set the data type and channels used for training. In this training, we use application/x-image content type that require individual images and LST file for data input.

Step 8: Train the model

Start training by calling the fit method in the estimator.

Inference

Step 9: Deploy the Model for Inference

Once we have trained the model, we want to use it to perform inference. In this case, we will be predicting keywords for the image.

First, let's deploy the model. You can deploy the created model by using the deploy method in the estimator.

Step 10: Evaluate the model

Next, we evaluate one of our images through the network for inference. The network outputs probabilities for each of the keywords that we are training on. As can be seen from this example, the network output is pretty good even with training for only 5 epochs.

After running the above lines of code, our notebook will display the image and print out the probabilities of each of the keywords that we are training on. From here, we can expand the amount of keywords that we train on, fine-tune our parameters, or add additional images to further improve the accuracy of our model.

Clean up

You can use the following command to delete the SageMaker model endpoint. The endpoint that is created above is persistent and would consume resources till it is deleted. It is good to delete the endpoint when it is not being used.

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.