TensorFlow Model Export

Introduction

After training a machine learning model with TensorFlow, the next crucial step is to export it in a format suitable for deployment. Model export is the process of converting your trained model into a format that can be used in production environments, whether it's a web application, mobile device, or cloud service.

In this guide, we'll explore different ways to export TensorFlow models, understand the various available formats, and learn the best practices for preparing your models for deployment.

Why Export Models?

Before diving into the how, let's understand why exporting models is important:

Deployment: Trained models need to be deployed to serve predictions in real-world applications.
Portability: Exported models can be used across different environments and platforms.
Optimization: Some export formats optimize models for specific hardware or reduce size.
Versioning: Export enables proper versioning and management of models.

SavedModel Format

The SavedModel format is TensorFlow's recommended format for exporting models. It contains a complete TensorFlow program, including weights and computation.

Basic Export to SavedModel

Here's how to export a simple model to the SavedModel format:

python
import tensorflow as tf
from tensorflow import keras

# Create a simple model
model = keras.Sequential([
    keras.layers.Dense(128, activation='relu', input_shape=(784,)),
    keras.layers.Dense(10)
])

# Compile the model
model.compile(
    optimizer='adam',
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

# Save the model in SavedModel format
export_path = './saved_models/mnist/1'  # Version number as a subdirectory
model.save(export_path)

print(f"Model saved to: {export_path}")

Output:

Model saved to: ./saved_models/mnist/1

Understanding the SavedModel Directory

After exporting, your SavedModel directory will contain:

assets/: Optional external files used by the model
variables/: Contains the model's weights
saved_model.pb: The model's architecture and computation graph

Let's inspect the SavedModel:

python
# Load the saved model back
loaded_model = tf.keras.models.load_model(export_path)

# Check its architecture
loaded_model.summary()

# You can also use the SavedModel CLI to inspect the model
# !saved_model_cli show --dir {export_path} --all

Exporting with Signatures

When exporting models, you can specify signatures that define the inputs and outputs. This is especially useful for serving.

python
# Create a simple model
model = keras.Sequential([
    keras.layers.Dense(10, activation='relu', input_shape=(5,)),
    keras.layers.Dense(1)
])

# Create a function that will be traced and exported
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 5], dtype=tf.float32)])
def serving_fn(inputs):
    return {'predictions': model(inputs)}

# Export the model with the signature
export_path = './saved_models/with_signature/1'
tf.saved_model.save(
    model, 
    export_path, 
    signatures={'serving_default': serving_fn}
)

print(f"Model with signature saved to: {export_path}")

TensorFlow Lite Conversion

TensorFlow Lite is designed for mobile and edge devices. It's a lightweight version of TensorFlow with a smaller binary size, fewer dependencies, and optimized for on-device ML.

Converting to TensorFlow Lite

python
import numpy as np

# Create and train a simple model
model = keras.Sequential([
    keras.layers.Dense(8, activation='relu', input_shape=(4,)),
    keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy')

# Train with some dummy data
x_train = np.random.random((100, 4))
y_train = np.random.randint(0, 2, (100, 1))
model.fit(x_train, y_train, epochs=3)

# Convert the model to TensorFlow Lite format
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# Save the TF Lite model
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

print("TensorFlow Lite model saved to 'model.tflite'")

Quantization for Size Reduction

To reduce the size of your TensorFlow Lite models, you can apply quantization:

python
# Convert with quantization
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # This enables quantization

# Representative dataset for full integer quantization (optional)
def representative_dataset():
    for i in range(10):
        yield [np.random.random((1, 4)).astype(np.float32)]
        
converter.representative_dataset = representative_dataset

# Convert the model
tflite_quant_model = converter.convert()

# Save the quantized model
with open('model_quantized.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print("Size of the original model:", len(tflite_model), "bytes")
print("Size of the quantized model:", len(tflite_quant_model), "bytes")

Output (example):

Size of the original model: 2512 bytes
Size of the quantized model: 1374 bytes

TensorFlow.js Export

TensorFlow.js lets you run machine learning models in browsers or with Node.js.

Converting for Web Deployment

First, install the tensorflowjs converter:

bash
pip install tensorflowjs

Then convert your model:

python
import tensorflowjs as tfjs

# First, save your model in SavedModel format
model.save('tfjs_model_tmp')

# Convert the model to TensorFlow.js format
tfjs.converters.save_keras_model(model, 'tfjs_model')

print("Model converted to TensorFlow.js format in 'tfjs_model' directory")

The converted model can be loaded in JavaScript like this:

javascript
// Example JavaScript code to load the model in a browser
async function loadModel() {
  const model = await tf.loadLayersModel('tfjs_model/model.json');
  console.log('Model loaded');
  return model;
}

// Make predictions
async function predict(model, input) {
  const tensor = tf.tensor2d([input]);
  const prediction = model.predict(tensor);
  return prediction.dataSync();
}

Real-World Example: Exporting an Image Classification Model

Let's export a pre-trained MobileNetV2 model in multiple formats:

python
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2

# Load a pre-trained model
base_model = MobileNetV2(weights='imagenet', include_top=True)

# Create a function to preprocess images and get predictions
@tf.function(input_signature=[tf.TensorSpec(shape=[None, None, None, 3], dtype=tf.uint8)])
def serve_function(input_img):
    # Preprocess the image
    img = tf.cast(input_img, tf.float32)
    img = tf.image.resize(img, (224, 224))
    img = tf.keras.applications.mobilenet_v2.preprocess_input(img)
    
    # Get predictions
    predictions = base_model(img)
    
    # Return predictions and decoded class names
    return {
        'predictions': predictions,
    }

# Export the model with preprocessing included
export_path = './mobilenet_savedmodel'
tf.saved_model.save(
    base_model, 
    export_path, 
    signatures={'serving_default': serve_function}
)

print(f"MobileNetV2 exported to: {export_path}")

# Convert to TensorFlow Lite
converter = tf.lite.TFLiteConverter.from_saved_model(export_path)
tflite_model = converter.convert()

# Save the TFLite model
with open('mobilenet.tflite', 'wb') as f:
    f.write(tflite_model)

print("MobileNetV2 also exported to TFLite format")

Best Practices for Model Export

Include Preprocessing: Bundle preprocessing steps with your model where possible.
Test After Export: Always verify that exported models produce the same outputs as the original.
Version Your Models: Use directory structures or naming conventions to keep track of model versions.
Document Input/Output Specs: Clearly document the expected input shapes, types, and preprocessing requirements.
Consider Optimization: Apply optimization techniques like quantization for edge deployment.
Check Dependencies: Ensure all required assets are included in the export.

Common Issues and Solutions

Model Size Too Large

If your exported model is too large:

Use quantization techniques
Prune unnecessary layers
Consider smaller architectures

Input Shape Issues

If you encounter input shape incompatibilities:

Use input_signature in tf.function decorators
Make your model work with dynamic shapes where possible
Document expected input shapes clearly

Missing Custom Objects

If your model uses custom layers or functions:

Use the custom_objects parameter when loading models
Make sure to register custom objects with Keras

python
# Example of handling custom objects
class CustomLayer(tf.keras.layers.Layer):
    # Layer implementation...
    pass

# When saving
model.save('model_with_custom_layer')

# When loading
custom_objects = {'CustomLayer': CustomLayer}
loaded_model = tf.keras.models.load_model('model_with_custom_layer', custom_objects=custom_objects)

Summary

In this guide, we've covered:

Exporting models using SavedModel format
Converting models to TensorFlow Lite for mobile and edge devices
Converting models to TensorFlow.js for web applications
Best practices for model export and common issues to avoid

Model export is a critical step in the machine learning lifecycle, bridging the gap between model development and deployment. By choosing the right export format and following best practices, you can ensure your models are ready for real-world applications.

Additional Resources

Exercises

Export a pre-trained ResNet50 model using the SavedModel format and add a preprocessing signature.
Convert the same model to TensorFlow Lite and compare the performance and size.
Create a custom model with a few layers, export it in all three formats (SavedModel, TFLite, TFJS), and verify that they produce the same outputs for test inputs.
Experiment with different quantization techniques on a model and measure the impact on size and accuracy.
Build a simple web application that uses your TensorFlow.js exported model to make predictions in the browser.

If you spot any mistakes on this website, please let me know at [email protected]. I’d greatly appreciate your feedback! :)

Introduction​

Why Export Models?​

SavedModel Format​

Basic Export to SavedModel​

Understanding the SavedModel Directory​

Exporting with Signatures​

TensorFlow Lite Conversion​

Converting to TensorFlow Lite​

Quantization for Size Reduction​

TensorFlow.js Export​

Converting for Web Deployment​

Real-World Example: Exporting an Image Classification Model​

Best Practices for Model Export​

Common Issues and Solutions​

Model Size Too Large​

Input Shape Issues​

Missing Custom Objects​

Summary​

Additional Resources​

Exercises​