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:
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.)