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
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
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
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
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
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:
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/
# Set via environment variable
export TFHUB_CACHE_DIR=/models/tfhub_cache
python train.py
Finding the Cache Path
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:
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
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
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
