Loading tensorflow models from Amazon S3 with Tensorflow Serving

In this article, I am going to show you how to store a Tensorflow model in a file, upload it to Amazon S3, and configure the Docker image of Tensorflow Serving to serve that model via REST API.

Table of Contents

Saving the model
Upload to S3
Configuring Tensorflow serving to use a model from S3

Saving the model

Before we start, we have to save a Tensorflow model in a file using the simple_save function.

I’m going to assume that you have already trained your model. We need to specify the output directory and make sure that such a location exists. When the target directory is ready, we can call the simple_save function.

import tensorflow as tf
from tensorflow import keras

export_path = "/home/user/some_location" #the absolute path to the target directory
model = ... #a trained Tensorflow model

tf.saved_model.simple_save(
    keras.backend.get_session(),
    export_path,
    inputs={'input_image': model.input},
    outputs={t.name:t for t in model.outputs})

Upload to S3

If you have Amazon command line tool installed, you can use the following command to upload the model to S3.

Note that you should replace /home/user/some_location with the location where you stored the model and tensorflow_models/example_model with your bucket name.

aws s3 cp /home/user/some_location s3://tensorflow_models/example_model --recursive

Configuring Tensorflow serving to use a model from S3

To serve the model, we are going to need five additional information. The most important part is your AWS account configuration: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and the name of the AWS region, which you used to store the file.

You should decide what the name of your model is. Tensorflow serving will use that name as a part of the REST endpoint URL.

We also need the address of the S3 endpoint. If you know the AWS region you use, you can find the S3 endpoint address in the documentation: https://docs.aws.amazon.com/general/latest/gr/rande.html#s3_region

Take a look at the following command and modify the aforementioned parameters.

docker run -p 8501:8501 \
    -e AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
    -e AWS_SECRET_ACCESS_KEY=$AWS_SECRET_ACCESS_KEY \
    -e MODEL_BASE_PATH=s3://tensorflow_models/example_model \
    -e MODEL_NAME=example_model \
    -e S3_ENDPOINT=s3.eu-central-1.amazonaws.com \
    -e AWS_REGION=eu-central-1 \
    -e TF_CPP_MIN_LOG_LEVEL=3 \
    -t tensorflow/serving

Now, you can run that Docker image, and it will load the model from Amazon S3.

Loading tensorflow models from Amazon S3 with Tensorflow Serving

Saving the model

Upload to S3

Configuring Tensorflow serving to use a model from S3

Pandas stack and unstack explained

A.I. in production: your next stylist is going to be a neural network

Loading tensorflow models from Amazon S3 with Tensorflow Serving

Saving the model

Upload to S3

Configuring Tensorflow serving to use a model from S3

Pandas stack and unstack explained

A.I. in production: your next stylist is going to be a neural network

Related Posts

How to deploy a Transformer-based model with custom preprocessing code to Sagemaker Endpoints using BentoML

Multimodel deployment in Sagemaker Endpoints

How to add custom preprocessing code to a Sagemaker Endpoint running a Tensorflow model