Introduction

After spending hours training a text classification model, the last thing you want is to lose it when your notebook session ends. TensorFlow provides several ways to persist a trained model so you can reload it for inference, share it with others, or deploy it to production. This article covers the main saving formats, how to handle custom objects, and how to export to TensorFlow Lite for mobile deployment.

Building a Sample Text Classification Model

Before discussing saving strategies, here is a minimal text classification model to use as a reference throughout the examples.

1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Embedding(input_dim=10000, output_dim=64),
5    tf.keras.layers.GlobalAveragePooling1D(),
6    tf.keras.layers.Dense(64, activation="relu"),
7    tf.keras.layers.Dropout(0.3),
8    tf.keras.layers.Dense(1, activation="sigmoid"),
9])
10
11model.compile(
12    optimizer="adam",
13    loss="binary_crossentropy",
14    metrics=["accuracy"],
15)
16
17# Assume training has been completed:
18# model.fit(train_data, train_labels, epochs=10, validation_split=0.2)

Saving with model.save() (SavedModel Format)

The recommended way to save a complete model is model.save(), which writes the architecture, weights, optimizer state, and training configuration to disk in TensorFlow's SavedModel format.

# Save the entire model as a SavedModel directory
model.save("saved_models/text_classifier")

This creates a directory containing a saved_model.pb file, a variables/ folder with weight checkpoints, and an assets/ folder. The SavedModel format is the default in TensorFlow 2.x and is the format expected by TensorFlow Serving, TensorFlow.js, and TensorFlow Lite converters.

To reload the model later:

1loaded_model = tf.keras.models.load_model("saved_models/text_classifier")
2
3# Verify it works
4predictions = loaded_model.predict(test_data)

The loaded model is identical to the original. You can resume training, run inference, or export it further.

Saving Weights Only with save_weights()

Sometimes you only need to save the learned parameters, not the full architecture. This is useful when you want to load weights into a model with a different configuration or when the architecture is defined in code and you want to version-control it separately.

1# Save only the weights
2model.save_weights("checkpoints/text_clf_weights.weights.h5")
3
4# Later, rebuild the same architecture and load weights
5new_model = tf.keras.Sequential([
6    tf.keras.layers.Embedding(input_dim=10000, output_dim=64),
7    tf.keras.layers.GlobalAveragePooling1D(),
8    tf.keras.layers.Dense(64, activation="relu"),
9    tf.keras.layers.Dropout(0.3),
10    tf.keras.layers.Dense(1, activation="sigmoid"),
11])
12new_model.load_weights("checkpoints/text_clf_weights.weights.h5")

The key requirement is that the architecture you load weights into must have exactly the same layer structure and shapes. If they differ, TensorFlow raises a shape mismatch error.

Handling Custom Objects

Text classification models often include custom layers, loss functions, or preprocessing steps. When you save a model that contains custom objects, TensorFlow needs to know how to reconstruct them at load time.

1class TextPreprocessor(tf.keras.layers.Layer):
2    def __init__(self, max_length=200, **kwargs):
3        super().__init__(**kwargs)
4        self.max_length = max_length
5
6    def call(self, inputs):
7        return inputs[:, :self.max_length]
8
9    def get_config(self):
10        config = super().get_config()
11        config.update({"max_length": self.max_length})
12        return config
13
14# Save normally
15model_with_custom = tf.keras.Sequential([
16    TextPreprocessor(max_length=200),
17    tf.keras.layers.Embedding(10000, 64),
18    tf.keras.layers.GlobalAveragePooling1D(),
19    tf.keras.layers.Dense(1, activation="sigmoid"),
20])
21model_with_custom.compile(optimizer="adam", loss="binary_crossentropy")
22model_with_custom.save("saved_models/custom_text_clf")
23
24# Load with custom_objects
25loaded = tf.keras.models.load_model(
26    "saved_models/custom_text_clf",
27    custom_objects={"TextPreprocessor": TextPreprocessor},
28)

The critical step is implementing get_config() in your custom layer. Without it, TensorFlow cannot serialize the layer's constructor arguments, and loading will fail. The custom_objects dictionary maps class names to their Python classes so the loader can reconstruct them.

Exporting to TensorFlow Lite

For deploying text classification models on mobile or edge devices, TensorFlow Lite provides a smaller, optimized format. You convert from a SavedModel to a .tflite file using the TFLiteConverter.

1# Convert the saved model to TFLite format
2converter = tf.lite.TFLiteConverter.from_saved_model("saved_models/text_classifier")
3tflite_model = converter.convert()
4
5# Write the TFLite model to disk
6with open("text_classifier.tflite", "wb") as f:
7    f.write(tflite_model)

You can also apply post-training quantization to reduce the model size further:

1converter.optimizations = [tf.lite.Optimize.DEFAULT]
2quantized_model = converter.convert()
3
4with open("text_classifier_quantized.tflite", "wb") as f:
5    f.write(quantized_model)

Quantization typically reduces model size by roughly 4x with minimal accuracy loss for text classification tasks.

Loading Back for Inference

Regardless of which format you saved in, loading for inference follows a consistent pattern.

1# From SavedModel
2model = tf.keras.models.load_model("saved_models/text_classifier")
3predictions = model.predict(new_texts)
4
5# From TFLite
6interpreter = tf.lite.Interpreter(model_path="text_classifier.tflite")
7interpreter.allocate_tensors()
8input_details = interpreter.get_input_details()
9output_details = interpreter.get_output_details()
10
11interpreter.set_tensor(input_details[0]["index"], input_data)
12interpreter.invoke()
13result = interpreter.get_tensor(output_details[0]["index"])

The SavedModel path is simpler and supports the full Keras API. The TFLite path requires manual tensor management but runs efficiently on resource-constrained devices.

Common Pitfalls

• Saving a model that has not been built yet (no input shape inferred), which causes model.save() to fail; always call model.build(input_shape) or run at least one forward pass first.
• Forgetting to implement get_config() on custom layers, making the model impossible to load from a SavedModel without workarounds.
• Using model.save("model.h5") with the legacy HDF5 format when your model contains custom objects or TF operations that HDF5 does not support; prefer the default SavedModel directory format.
• Attempting to load weights into a model with a different architecture than the one used during saving, which produces shape mismatch errors that can be confusing.
• Not testing the loaded model with sample input immediately after saving, which delays catching serialization bugs until deployment.

Summary

• Use model.save() with the SavedModel format to persist the complete model (architecture, weights, optimizer state, and training config).
• Use save_weights() when you want to version architecture in code and only persist learned parameters.
• Implement get_config() on all custom layers and pass custom_objects when loading models that use them.
• Convert to TensorFlow Lite with TFLiteConverter for mobile and edge deployment, optionally applying quantization for smaller file sizes.
• Always test the loaded model with sample input immediately after saving to verify serialization correctness.