Build an Object Detection Model in TensorFlow: Model Training and Evaluation

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.

Background

This notebook is one of a sequence of notebooks that show you how to use various SageMaker functionalities to build, train, and test the object detection model, including data pre-processing steps like ingestion, cleaning and processing, training, and test the model. There are two parts of the demo:

1. Overview and Data Preparation.- you will preprocess the data, then create a json file from the cleaned data. By the end of part 1, you will have a complete data set that contains all features used on Object selection to be ingested by a data loader in TensorFlow using 'TFRecords'.
2. Data loader creation and Model Training (current notebook).- you will use the data set built from part 1 to create a data loader for TensorFlow using Keras CV, train the model and then test the model predictability with the test data.

Content

• Overview
• Using TensorFlow Data loaders with 'TFRecords'
• Model Selection
• Training Object Detection Models with SageMaker and TensorFlow

Overview

What is Object Detection, and why is it important?

Object detection refers to detecting instances of objects from certain classes in images or videos. It allows for multiple objects to be detected and localized in an image. Object detection is commonly used in applications such as self-driving cars, face detection, video surveillance, etc.

Use Cases for Object Detection

Some common use cases of object detection include:

• Self driving cars - detect pedestrians, cars, traffic signs, etc.
• Face detection - detect faces in images and videos for applications like security and tagging people in images.
• Video surveillance - detect suspicious activities or objects.
• Medical imaging - detect anomalies, tumors, etc. in medical scans.
• Retail - detect objects on shelves for inventory management.

Define the Machine Learning Problem

Object detection can be formulated as a supervised machine learning problem:

• Given a set of labelled images containing objects from certain classes, train a model to detect the presence and location of those objects in new images.

• The model needs to identify the class of objects present and draw bounding boxes around them indicating their locations.

Data Requirements

• Large dataset of images with object annotation - Object locations are annotated using bounding boxes around them.

• Variety of images - Objects captured under different conditions of illumination, scales, occlusion, viewpoints etc.

Challenges

• Data annotation - Time consuming and expensive process.

• Class imbalance - Models tend to perform better for classes with more examples.

• Viewpoint variation - Objects look different from different angles and viewpoints.

• Background clutter - Objects may blend with their surroundings.

• Small objects - Harder to detect smaller objects.

• Occlusion - Objects hidden behind other things are tougher to detect.

Parameters

The following lists configurable parameters that are used throughout the whole notebook.

Using TensorFlow Data loaders with 'TFRecords'

Loading data from S3

As a first step, we need to load the 'TFRecord' files that were generated during the preprocessing job and saved to S3. We can use the TensorFlow IO functions to stream data directly from S3 without needing to download the files locally.

The code starts by loading variables from a previous notebook using the '%store -r magic command':

It generates a class_mapping dictionary that maps class IDs to their corresponding class names. This mapping will be used later for visualization purposes.

This code cell creates a SageMaker session and gets the default S3 bucket associated with the session. The S3 client is also instantiated to interact with the S3 API. The cell defines the S3 prefixes where the preprocessed training and validation data are stored in 'TFRecord' format. Finally, it lists the objects (files) in the S3 bucket under the specified prefixes for training and validation data, using the list_objects_v2 method of the S3 client. This step is necessary to retrieve the paths of the 'TFRecord' files, which will be used for loading the data.

The provided code retrieves the data and prepares the datasets for training and validation. It accomplishes this by generating lists of local file paths for the 'TFRecord' files and subsequently loading those files into TensorFlow datasets. These prepared datasets can then be utilized for training or evaluating the object detection model.

Using data loaders and parsing 'TFRecords'

Next we can create a data loader to parse the 'TFRecord' examples and create batches of images and labels for visualization.

This cell sets up the necessary functions and imports for working with 'TFRecord' data in the context of object detection. The 'parse_tfrecord_fn' function defines how to parse the 'TFRecord' file format, which stores the image data and bounding box annotations. It reads the features from the 'TFRecord' file, decodes the image data, and creates a dictionary with the image and bounding box information. The prepare_sample function is a helper function that formats the parsed data into the expected format for the object detection model. The 'plot_boxes_tfrecords' function is a utility for visualizing the bounding boxes on the images.

This code prepares the 'TFRecord' dataset to test an object detection model. It loads the dataset, applies necessary preprocessing, shuffles the data, and retrieves a sample for visualization.

Data Augmentation

One major benefit of using TensorFlow data loaders is that we can easily apply data augmentation. This helps prevent overfitting and improves the robustness of the model.

Some common augmentation techniques for object detection include:

• Random horizontal/vertical flipping
• Random cropping
• Color jittering
• Adding noise

These can be implemented using the Keras CV layers for computer vision use cases:

Adding augmentation during training helps prevent overfitting and makes the model more robust to variations in input images.

Model Selection

You will need to select an appropriate model architecture for your object detection task. When choosing a model, there are several factors to consider:

• The type of objects you want to detect - Are they general everyday objects, or more specialized categories like faces or text? Simpler architectures like SSD and YOLO work well for detecting common objects, while more complex models like Mask R-CNN may be better for niche categories.

• Model size and speed - Larger models like 'RetinaNet' will be more accurate but slower, while smaller models like MobileNet will be faster but less accurate. Choose a model size that fits your speed and accuracy needs.

• Amount of training data - If you have a large dataset, you can train bigger models with more parameters. With fewer data, stick to smaller models to avoid overfitting.

• Inference speed - Some models like 'MobileNet' are optimized specifically for fast inference after training. Prioritize this if you need to run detection very quickly.

• Built-in vs custom models - Many pre-made model architectures like 'Faster R-CNN' are available. But you can also build custom models better tailored to your specific objects.

A good starting point is to evaluate pre-trained models like 'Faster R-CNN' and 'SSD' (Single Shot Detector) that are available in model zoos. 'Faster R-CNN' with a 'ResNet-50' backbone offers a good balance of accuracy and speed for this dataset. 'SSD' is faster but slightly less accurate, so you may want to try different backbone architectures like 'ResNet', 'MobileNet' and 'EfficientNet' to find the right tradeoff.

To simplify model development, we will leverage the pre-trained object detection models available in Keras CV. Keras CV provides reference implementations and pre-trained weights for state-of-the-art computer vision models. We can quickly test training and inference for object detection by using a model like 'RetinaNet' or 'EfficientNet', initialized with weights pre-trained on COCO or other datasets. By taking advantage of these pre-trained models in Keras CV, we can prototype and experiment with minimal code and set up time. This allows us to focus on customizing and optimizing the model for our specific use case.

How to Load Pretrained models

Loading Pretrained models

This cell loads a pre-trained RetinaNet object detection model based on the ResNet50 architecture and the Pascal VOC dataset. The from_preset function is used to load the pre-trained weights and architecture. The bounding_box_format parameter specifies the format of the bounding box coordinates, which is 'xyxy' in this case. The prediction_decoder parameter is set to the 'NonMaxSuppression' layer initialized in the previous cell, which will be used to filter out overlapping bounding boxes during inference. The load_weights parameter is set to True to load the pre-trained weights along with the model architecture.

This code creates a resizing layer for inference. The 'keras_cv.layers.Resizing' layer is used to resize the input images to a fixed size of 640x640 pixels. The bounding_box_format parameter specifies the format of the bounding box coordinates, where 'xyxy' means that the coordinates are in the format of [x_min, y_min, x_max, y_max]. The pad_to_aspect_ratio parameter ensures that the aspect ratio of the images is preserved during resizing by adding padding if necessary.

Prediction test

To validate the performance of a pretrained model, we can run a prediction test.

Training Object Detection Models with SageMaker and TensorFlow

To train object detection models on SageMaker, we first need to configure a SageMaker training job. The key components are:

• Choosing an estimator
  • We can use the TensorFlow estimator to leverage the Keras API and pretrained models like RetinaNet.
• Selecting an instance type
  • GPU instances like ml.p3.2xlarge are best suited for training convolutional neural networks.
• Configuring the training script
  • This sets up the model architecture, loads pretrained weights, and defines the training loop.
• Specifying the training image
  • We can use a TensorFlow image from the SageMaker registry.
• Setting hyperparameters
  • Learning rate, batch size, and epochs are key hyperparameters to tune.

For model evaluation, we need to choose appropriate metrics like precision, recall, and 'mAP'. Since object detection involves classifying many bounding boxes, metrics that account for class imbalance like F1 score are also useful. The pretrained models in Keras CV combined with SageMaker's managed training provide an optimized environment for iterating on object detection models. We can efficiently improve accuracy by tuning hyperparameters and leverage SageMaker infrastructure for scalable distributed training.

How to create a training job in Sagemaker

Loading Preprocessed Data from S3

A key advantage of using SageMaker for model training is it can directly access data stored in S3 buckets. After preprocessing our dataset in a previous step, we staged the output in an S3 location. The SageMaker TensorFlow estimator handles loading this data from S3 into our training script. We simply specify the S3 path when creating the TensorFlow estimator.

If you want to see how the train and validation 'TFRecords' datasets are created in detail, look at Build an Object Detection Model on Tensorflow and SageMaker: Overview and Data Preparation.

Initialize Model Hyperparameters

Define SageMaker estimator

This code sets up a TensorFlow estimator for training an object detection model using SageMaker.

Create Training Job

The next cell starts the training job for the object detection model. The fit method of the estimator object is used to initiate the training process.

How to load models from training jobs to use them locally

Load last training job metadata

This cell imports the necessary AWS SDK for Python (Boto3) and retrieves the name of the last training job from the SageMaker service. It uses the list_training_jobs API call to get a list of training jobs sorted by creation time in descending order, and takes the first result (the most recent training job that was completed).

This cell calls the describe_training_job API to retrieve the details of the training job with the given name. Finally, it prints the S3 URI of the model artifacts generated by the training job.

Download model artifacts from S3 bucket for local testing

The code of the next cells downloads a pre-trained model from an Amazon S3 bucket and extracts it locally.

Load model locally

This code loads the pre-trained model locally.

Load validation dataset

This code is preparing a validation dataset for a machine learning model that performs object detection. It loads and preprocesses the validation data from 'TFRecord' files, shuffles and batches the data, resizes the input images to a fixed size, and converts the input data to a format suitable for the model. The preprocessed validation dataset is then ready for evaluating the model's performance on unseen data.

Local prediction testing

The 'keras_cv.layers.MultiClassNonMaxSuppression' layer is used to create a prediction decoder for the object detection model. This layer performs non-maximum suppression on the raw output of the model, which helps to remove duplicate or overlapping bounding box predictions.

This function, visualize_detections, is used to visualize the object detection results of a trained model on a sample dataset

Create metrics for trained model

This code is setting up an instance of the 'BoxCOCOMetrics' class from the keras_cv.metrics module. This class is used to calculate various evaluation metrics for object detection models, specifically when working with the COCO (Common Objects in Context) dataset format.

The validation metrics computed in this code can be used to evaluate the performance of the trained model and potentially fine-tune it if necessary.

This code is useful for evaluating the performance of an object detection model on a validation dataset, as it calculates various metrics such as precision, recall, and mean average precision ('mAP') based on the true and predicted labels. These metrics can be used to assess the model's accuracy and make necessary adjustments or improvements.

Note: To achieve good Mean Average Precision ('mAP') values, you may need to run the training for more epochs and perform hyperparameter tuning. This code is just an exercise, and further optimization might be required for real-world object detection tasks.

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.