Migrating from TF1 to TF2 with TensorFlow Hub

This page explains how to keep using TensorFlow Hub while migrating your TensorFlow code from TensorFlow 1 to TensorFlow 2. It complements TensorFlow's general migration guide.

For TF2, TF Hub has switched away from the legacy hub.Module API for building a tf.compat.v1.Graph like tf.contrib.v1.layers do. Instead, there is now a hub.KerasLayer for use alongside other Keras layers for building a tf.keras.Model (typically in TF2's new eager execution environment) and its underlying hub.load() method for low-level TensorFlow code.

The hub.Module API remains available in the tensorflow_hub library for use in TF1 and in the TF1 compatibility mode of TF2. It can only load models in the TF1 Hub format.

The new API of hub.load() and hub.KerasLayer works for TensorFlow 1.15 (in eager and graph mode) and in TensorFlow 2. This new API can load the new TF2 SavedModel assets, and, with the restrictions laid out in the model compatibility guide, the legacy models in TF1 Hub format.

In general, it is recommended to use new API wherever possible.

Summary of the new API

hub.load() is the new low-level function to load a SavedModel from TensorFlow Hub (or compatible services). It wraps TF2's tf.saved_model.load(); TensorFlow's SavedModel Guide describes what you can do with the result.

m = hub.load(handle)
outputs = m(inputs)

The hub.KerasLayer class calls hub.load() and adapts the result for use in Keras alongside other Keras layers. (It may even be a convenient wrapper for loaded SavedModels used in other ways.)

model = tf.keras.Sequential([
    hub.KerasLayer(handle),
    ...])

Many tutorials show these APIs in action. Here are some examples:

Text classification example notebook
Image classification example notebook

Using the new API in Estimator training

If you use a TF2 SavedModel in an Estimator for training with parameter servers (or otherwise in a TF1 Session with variables placed on remote devices), you need to set experimental.share_cluster_devices_in_session in the tf.Session's ConfigProto, or else you will get an error like "Assigned device '/job:ps/replica:0/task:0/device:CPU:0' does not match any device."

The necessary option can be set like

session_config = tf.compat.v1.ConfigProto()
session_config.experimental.share_cluster_devices_in_session = True
run_config = tf.estimator.RunConfig(..., session_config=session_config)
estimator = tf.estimator.Estimator(..., config=run_config)

Starting with TF2.2, this option is no longer experimental, and the .experimental piece can be dropped.

Loading legacy models in TF1 Hub format

It can happen that a new TF2 SavedModel is not yet available for your use-case and you need to load an legacy model in TF1 Hub format. Starting in tensorflow_hub release 0.7, you can use legacy model in TF1 Hub format together with hub.KerasLayer as shown below:

m = hub.KerasLayer(handle)
tensor_out = m(tensor_in)

Additionally KerasLayer exposes the ability to specify tags, signature, output_key and signature_outputs_as_dict for more specific usages of legacy models in TF1 Hub format and legacy SavedModels.

For more information on TF1 Hub format compatibility see the model compatibility guide.

Using lower level APIs

Legacy TF1 Hub format models can be loaded via tf.saved_model.load. Instead of

# DEPRECATED: TensorFlow 1
m = hub.Module(handle, tags={"foo", "bar"})
tensors_out_dict = m(dict(x1=..., x2=...), signature="sig", as_dict=True)

it is recommended to use:

# TensorFlow 2
m = hub.load(path, tags={"foo", "bar"})
tensors_out_dict = m.signatures["sig"](x1=..., x2=...)

In these examples m.signatures is a dict of TensorFlow concrete functions keyed by signature names. Calling such a function computes all its outputs, even if unused. (This is different from the lazy evaluation of TF1's graph mode.)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!