Deploy a TensorFlow trained image classification model to AWS DeepLens

We are very excited to announce that you can deploy your computer vision model trained using TensorFlow (version 1.4) to AWS DeepLens. Head pose detection is part of the AWS DeepLens sample projects. In this blog post, we will show you how to train a model from scratch using a P2 training instance of Amazon SageMaker. We will use a ResNet-50 model and save the trained model in the “frozen” protobuff format. Although TensorFlow offers a variety of data formats to save a model graph (.such as checkpoint, .ckpt-XXX.meta, .ckpt_XXX.index, .ckpt_XXX.data-00000-of-00001, .pbtxt, optimized protobuff, frozen protobuff, etc.), the AWS DeepLens model optimizer supports the frozen protobuff format.

We use the same Prima head pose dataset that we previously described in Deploy Gluon models to AWS DeepLens using a simple Python API.

Create an Amazon S3 bucket

Just like we did in our last blog post, we are first going to create an Amazon S3 bucket using the Amazon S3 console. In this example, we are going to name the S3 bucket “deeplens-sagemaker-0001” hosted in the N. Virginia (US East 1) AWS Region. (If you want to deploy your trained model artifacts straight into AWS DeepLens, the Region must be N. Virginia (US East 1)).

Inside the bucket, we have a folder named “headpose.” Inside the “headpose” folder, we have 4 sub-folders named “TFartifacts,” “customTFcodes,” “datasets,” and “testIMs.”

You are going to host the head pose dataset from the previous blog post linked to earlier (HeadPoseData_trn_test_x15_py2.pkl) in the datasets folder.

That is it for the preparation.

Amazon SageMaker notebook

Now, let’s launch Amazon SageMaker. After you open your Amazon SageMaker notebook, upload our sample notebook, TensorFlow ResNet script, and entry point Python script (tensorflow_resnet_headpose_for_deeplens.ipynb, resnet_model_headpose.py, and resnet_headpose, respectively). These scripts are modified from Amazon SageMaker sample scripts.

After you add the notebook and Python scripts, there are only three steps for you to run the training.

First, specify your S3 bucket name in the sample Amazon SageMaker notebook (tensorflow_resnet_headpose_for_deeplens.ipynb). In this part, you also specify other folders inside your S3 bucket such as the “headpose” folder as well as the “TFartifacts” and “customTFcodes” folders underneath it.

import os
import sagemaker
from sagemaker import get_execution_role

sagemaker_session = sagemaker.Session()

s3_bucket = 'deeplens-sagemaker-0001'
headpose_folder = 'headpose'

#Bucket location to save your custom code in tar.gz format.
custom_code_folder = 'customTFcodes'
custom_code_upload_location = 's3://{}/{}/{}'.format(s3_bucket, headpose_folder, custom_code_folder)

#Bucket location where results of model training are saved.
model_artifacts_folder = 'TFartifacts'
model_artifacts_location = 's3://{}/{}/{}'.format(s3_bucket, headpose_folder, model_artifacts_folder)

#IAM execution role that gives SageMaker access to resources in your AWS account.
#We can use the SageMaker Python SDK to get the role from our notebook environment. 

role = get_execution_role()

Second, specify the training instance and other parameters in the TensorFlow object. In this example, we use a ml.p2.xlarge instance for the training.

from sagemaker.tensorflow import TensorFlow

source_dir = os.path.join(os.getcwd())

estimator = TensorFlow(entry_point='resnet_headpose.py',
                       framework_version = 1.4,
                       source_dir=source_dir,
                       role=role,
                       training_steps=20000, evaluation_steps=500,
                       train_instance_count=1, 
                       base_job_name='deeplens-TF-headpose',
                       output_path=model_artifacts_location,
                       code_location=custom_code_upload_location,
                       train_instance_type='ml.p2.xlarge',
                       train_max_run = 432000,
                       train_volume_size=100)


# Head-pose dataset "HeadPoseData_trn_test_x15_py2.pkl" is in the following S3 folder. 
dataset_location = 's3://{}/{}/datasets'.format(s3_bucket, headpose_folder)

AWS DeepLens currently supports TensorFlow version 1.4 (as of August 9, 2018). Amazon SageMaker helps you with version control by simply stating framework_version = 1.4.

Finally, we run the training by calling the “.fit” method.

estimator.fit(dataset_location)

After the training is over, you will find a set of trained TensorFlow model artifacts (model.tar.gz) inside an output folder in the TFartifacts folder of your S3 bucket.

Make a frozen protobuff file for AWS DeepLens

We are going to generate a frozen protobuff file from model.tar.gz, which we just made. In this tutorial, we use the TensorFlow Python API in the same Amazon SageMaker notebook.

First, we download the compressed file containing the model artifacts into the Amazon SageMaker notebook local directory.

import boto3
s3 = boto3.resource('s3')
key = '{}/{}/{}/output/model.tar.gz'.format(headpose_folder, model_artifacts_folder,estimator.latest_training_job.name)
print(key)
s3.Bucket(s3_bucket).download_file(key,'model.tar.gz')

After you decompress the file, you will find three separate files, saved_model.pb, variables/variables.index, and variables/variables.data-00000-of-00001 inside the export/Servo/{Assigned by Amazon SageMaker} directory.

import glob
model_dir = glob.glob('export/*/*')
print(model_dir)

Here is the code to freeze the graph and save it in the frozen protobuff format.

import tensorflow as tf
from tensorflow.python.tools import optimize_for_inference_lib
def freeze_graph(model_dir, output_node_names):
    """Extract the sub graph defined by the output nodes and convert 
    all its variables into constant 
    Args:
        model_dir: the root folder containing the checkpoint state file
        output_node_names: a string, containing all the output node's names, 
                            comma separated
    """
    
    # We start a session using a temporary fresh Graph
    with tf.Session(graph=tf.Graph()) as sess:
        # We import the meta graph in the current default Graph
        tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], model_dir)

        # We use a built-in TF helper to export variables to constants
        input_graph_def = tf.graph_util.convert_variables_to_constants(
            sess, # The session is used to retrieve the weights
            tf.get_default_graph().as_graph_def(), # The graph_def is used to retrieve the nodes 
            output_node_names.split(",") # The output node names are used to select the usefull nodes
        ) 

    # We generate the inference graph_def
    output_graph_def = optimize_for_inference_lib.optimize_for_inference(tf.graph_util.remove_training_nodes(input_graph_def),
                                                                         ['Const_1'], # an array of the input node(s)
                                                                         output_node_names.split(","), # an array of output nodes
                                                                         tf.float32.as_datatype_enum)
    # Finally we serialize and dump the output graph_def to the filesystem
    with tf.gfile.GFile('frozen_model.pb', "wb") as f:
            f.write(output_graph_def.SerializeToString())
    print("tf magic!")

We use the tf.graph_util.convert_variables_to_constants API to freeze the graph, then tf.graph_util.remove_training_nodes to remove all unnecessary nodes. Then, we use optimize_for_inference_lib.optimize_for_inference to generate the inference graph_def. Finally, we serialize and save the file as a protobuff.

freeze_graph(model_dir[-1], 'softmax_tensor')

It should be noted that we knew the names of input and output nodes in advance (that is, ‘Const_1’ and ‘softmax_tensor’). It is important to look into the graph using TensorBoard. In addition, it is also a good practice to name every single layer in the model script (resnet_model_headpose.py) to avoid the unexpected stray nodes/namespaces inside the graph.

After the frozen protobuff file, frozen_mode.pb is generated, we are going to put it back into the output folder inside the TFartifacts folder of the S3 bucket.

data = open('frozen_model.pb', "rb")
key = '{}/{}/{}/output/frozen_model.pb'.format(headpose_folder, model_artifacts_folder,estimator.latest_training_job.name)
s3.Bucket(s3_bucket).put_object(Key=key, Body=data)

You will find frozen_model.pb in your S3 bucket.

You can immediately deploy the model into your AWS DeepLens device.

The model is now ready. AWS DeepLens requires an AWS Lambda function to run the model. The Lambda function comes with one of the sample projects of AWS DeepLens.

If you want to learn how to write your own Lambda functions for AWS DeepLens, take a look at this blog post.

Conclusion

In this blog post, we trained a head-pose estimator ResNet-50 model in TensorFlow on Amazon SageMaker. Then we processed the trained model artifact file so that we can deploy it to the AWS DeepLens device. Now, you can develop your own AWS DeepLens model using TensorFlow on Amazon SageMaker.

About the Authors

Tatsuya Arai Ph.D. is a biomedical engineer turned deep learning data scientist on the Amazon Machine Learning Solutions Lab team. He believes in the true democratization of AI and that the power of AI shouldn’t be exclusive to computer scientists or mathematicians.

Eddie Calleja is a Software Development Engineer for AWS Deep Learning. He is one of the developers of the DeepLens device. As a former physicist he spends his spare time thinking about applying AI techniques to modern day physics problems.

Jyothi Nookula is a Senior Product Manager for AWS DeepLens. She loves to build products that delight her customers. In her spare time, she loves to paint and host charity fund raisers for her art exhibitions.