How to Deploy TensorFlow Model Endpoints on AWS

Select Language:

If you’re working with TensorFlow models in SageMaker and want to deploy your trained model as an endpoint, but encounter issues with model saving and deployment, here’s a simple step-by-step guide to get it done smoothly.

First, it’s important to understand that SageMaker’s TensorFlow container expects models to be saved in a specific format called the SavedModel directory structure. When you save your model using Keras’s save_model function with a .h5 extension, it saves the model in HDF5 format, which the TensorFlow container does not recognize for deployment by default.

To fix this, after training your model, you need to save it as a TensorFlow SavedModel directory, not as a .h5 file. Here’s how you can do it:

Save Your Model Correctly:
Instead of saving the model as a .h5 file, save it in the SavedModel format by adjusting your code like this:

python
model_path = os.path.join(args.model_dir, “saved_model”)
model.model.save(model_path, save_format=’tf’)

This creates a directory with the SavedModel structure that SageMaker expects.

Update the Model Data Path:
When deploying, specify the path to this SavedModel directory. If you’re uploading your model manually, zip this directory to a .tar.gz file before providing it to SageMaker, as it expects models in archive format.
Deploy the Model:
Use the TensorFlowModel class with the correct model data location. For example:

python
from sagemaker.tensorflow.model import TensorFlowModel

model = TensorFlowModel(
model_data=’s3://your-bucket/path-to-saved-model.tar.gz’,
role=role,
framework_version=’2.12′,
entry_point=’your_inference_script.py’ # this script handles inference
)

predictor = model.deploy(
initial_instance_count=1,
instance_type=’ml.m5.2xlarge’,
endpoint_name=’your-endpoint’
)

Make sure that the archive uploaded contains the SavedModel directory.

Simplify Your Workflow:
Instead of manually creating .gz files, you can use SageMaker’s SDK to upload your SavedModel directory directly to S3, and it will handle packaging for you.

By saving your model in the SavedModel format and properly packaging it, SageMaker can successfully deploy your model as an endpoint. This approach is straightforward and avoids the complexity of manually zipping files.

If you follow these steps, you’ll be able to deploy your model smoothly without the “no SavedModel bundles found” error.