Practice Now
tensorflow
machine-learning
tensorflow-hub
model-saving
python

How to save/load a tensorflow hub module to/from a custom path?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
TensorFlow & KerasMachine LearningPython

Introduction

To save a TensorFlow Hub module to a custom path, load it with hub.load() or hub.KerasLayer(), then use tf.saved_model.save() to export it to your directory. To load from a custom path, pass the local directory path to hub.load() or hub.KerasLayer() instead of a URL. You can also set the TFHUB_CACHE_DIR environment variable to control where TF Hub caches downloaded modules. This is essential for offline deployments, air-gapped environments, and reproducible model serving.

Saving a TF Hub Module to a Custom Path

Method 1: tf.saved_model.save

python
1import tensorflow as tf
2import tensorflow_hub as hub
3
4# Load from TF Hub URL
5module_url = "https://tfhub.dev/google/universal-sentence-encoder/4"
6model = hub.load(module_url)
7
8# Save to custom path
9custom_path = "/models/sentence_encoder"
10tf.saved_model.save(model, custom_path)
11
12print(f"Saved to {custom_path}")
13# Verify
14loaded = tf.saved_model.load(custom_path)

Method 2: Save as Part of a Keras Model

python
1import tensorflow as tf
2import tensorflow_hub as hub
3
4# Build a Keras model with a Hub layer
5model = tf.keras.Sequential([
6    hub.KerasLayer("https://tfhub.dev/google/nnlm-en-dim128/2",
7                   input_shape=[], dtype=tf.string),
8    tf.keras.layers.Dense(64, activation='relu'),
9    tf.keras.layers.Dense(1, activation='sigmoid')
10])
11
12model.compile(optimizer='adam', loss='binary_crossentropy')
13
14# Save the entire model (includes the Hub layer)
15model.save("/models/text_classifier")
16
17# Load later
18loaded_model = tf.keras.models.load_model("/models/text_classifier",
19    custom_objects={'KerasLayer': hub.KerasLayer})

Loading from a Custom Path

Method 1: hub.load with Local Path

python
1import tensorflow_hub as hub
2
3# Load from local directory instead of URL
4model = hub.load("/models/sentence_encoder")
5
6# Use the model
7embeddings = model(["Hello world", "TensorFlow Hub"])
8print(embeddings.shape)  # (2, 512)

Method 2: hub.KerasLayer with Local Path

python
1import tensorflow as tf
2import tensorflow_hub as hub
3
4# Use local path in KerasLayer
5layer = hub.KerasLayer("/models/sentence_encoder",
6                       trainable=False,
7                       input_shape=[],
8                       dtype=tf.string)
9
10model = tf.keras.Sequential([
11    layer,
12    tf.keras.layers.Dense(10, activation='softmax')
13])

Method 3: tf.saved_model.load

python
1import tensorflow as tf
2
3# Direct TensorFlow loading (no hub dependency needed)
4model = tf.saved_model.load("/models/sentence_encoder")
5
6# Call the default signature
7result = model(tf.constant(["test input"]))

Using TFHUB_CACHE_DIR

TF Hub caches downloaded modules. Control the cache location:

python
1import os
2
3# Set cache directory before importing hub
4os.environ['TFHUB_CACHE_DIR'] = '/models/tfhub_cache'
5
6import tensorflow_hub as hub
7
8# First call downloads and caches
9model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
10
11# Subsequent calls load from cache
12model2 = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
13# No download — loaded from /models/tfhub_cache/
bash
# Set via environment variable
export TFHUB_CACHE_DIR=/models/tfhub_cache
python train.py

Finding the Cache Path

python
1import tensorflow_hub as hub
2
3# Default cache location
4# Linux/Mac: /tmp/tfhub_modules/
5# Can also be: ~/.cache/tfhub_modules/
6
7# List cached modules
8import os
9cache_dir = os.environ.get('TFHUB_CACHE_DIR', '/tmp/tfhub_modules')
10if os.path.exists(cache_dir):
11    for item in os.listdir(cache_dir):
12        print(item)

Offline Deployment

For air-gapped or production environments without internet access:

python
1# Step 1: Download on a machine with internet
2import tensorflow_hub as hub
3import tensorflow as tf
4
5model = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
6tf.saved_model.save(model, "/export/sentence_encoder")
7
8# Step 2: Copy /export/sentence_encoder to the offline machine
9
10# Step 3: Load on the offline machine
11model = tf.saved_model.load("/path/to/sentence_encoder")
12embeddings = model(["Hello"])

Docker Deployment

dockerfile
1FROM tensorflow/tensorflow:latest
2
3# Pre-download TF Hub modules during build
4ENV TFHUB_CACHE_DIR=/models/tfhub_cache
5
6RUN python -c "\
7import os; \
8os.environ['TFHUB_CACHE_DIR'] = '/models/tfhub_cache'; \
9import tensorflow_hub as hub; \
10hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')"
11
12COPY app.py .
13CMD ["python", "app.py"]

Saving with Specific Signatures

python
1import tensorflow as tf
2import tensorflow_hub as hub
3
4module = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
5
6# Save with explicit signatures for serving
7class ServingModule(tf.Module):
8    def __init__(self, hub_module):
9        super().__init__()
10        self.module = hub_module
11
12    @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
13    def encode(self, texts):
14        return self.module(texts)
15
16serving = ServingModule(module)
17tf.saved_model.save(serving, "/models/serving_encoder",
18                    signatures={'encode': serving.encode})
19
20# Load and use the signature
21loaded = tf.saved_model.load("/models/serving_encoder")
22result = loaded.signatures['encode'](tf.constant(["hello"]))

Common Pitfalls

  • Setting TFHUB_CACHE_DIR after import tensorflow_hub: The cache directory is read when the module is first imported. Set the environment variable before importing tensorflow_hub, or it uses the default location.
  • Loading a Hub module saved with a different TF version: SavedModels may use ops not available in older TensorFlow versions. Always use the same major TensorFlow version for saving and loading, or use tf.compat.v1 for backward compatibility.
  • Forgetting custom_objects when loading a Keras model with Hub layers: tf.keras.models.load_model does not know about hub.KerasLayer by default. Pass custom_objects={'KerasLayer': hub.KerasLayer} to avoid ValueError: Unknown layer.
  • Assuming the cached module path is stable: TF Hub generates hashed directory names in the cache. Do not hardcode cache paths — use tf.saved_model.save() to export to a known, stable path.
  • Not checking disk space for large models: Some TF Hub modules (e.g., BERT, T5) are several gigabytes. Ensure the cache directory and export path have sufficient disk space, especially in Docker containers with limited storage.

Summary

  • Save with tf.saved_model.save(hub.load(url), "/path/to/model")
  • Load with hub.load("/path/to/model") or tf.saved_model.load("/path/to/model")
  • Set TFHUB_CACHE_DIR environment variable to control the download cache location
  • For offline deployments, export the model on a connected machine and copy the saved directory
  • Pass custom_objects={'KerasLayer': hub.KerasLayer} when loading Keras models with Hub layers
Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.