Thanks for tuning in to Google I/O. View all sessions on demandWatch on demand

hub.Module

Part of a TensorFlow 1 model that can be transferred between models.

hub.Module(
    spec, trainable=False, name='module', tags=None
)

Used in the notebooks

Used in the tutorials

A Module represents a part of a TensorFlow graph that can be exported to disk (based on the SavedModel format) and later re-loaded. A Module has a defined interface that allows it to be used in a replaceable way, with little or no knowledge of its internals and its serialization format. Example:

m = hub.Module("/tmp/text-embedding")
embeddings = m(sentences)

The module to instantiate is defined by its spec (a ModuleSpec or a path where to load it from) which contains the module weights, assets and signatures.

During instantiation the Module adds the state (e.g. variables and tables ops) to the current graph. Afterwards, the method __call__() allows to apply the module signatures multiple times, which adds ops for the computation.

A Module may provide different variants of its graph for different purposes (say, training or serving, which may behave differently, e.g., for batch normalization). Graph variants are identified by sets of string-valued tags. The graph variant used to create a module that is exported must define all the variables needed by any other graph variant that is subsequently used.

To make it possible to easily replace a module with another, they all assume that they will be used with common TensorFlow conventions such as session initialization and restore, use of collections for variables, regularization losses and updates, etc.

THIS FUNCTION IS DEPRECATED.

spec A ModuleSpec defining the Module to instantiate or a path where to load a ModuleSpec from via load_module_spec.
trainable whether the Module is trainable. If False, no variables are added to TRAINABLE_VARIABLES collection, and no tensors are added to REGULARIZATION_LOSSES collection.
name A string, the variable scope name under which to create the Module. It will be uniquified and the equivalent name scope must be unused.
tags A set of strings specifying the graph variant to use.

RuntimeError explaning the reason why it failed to instantiate the Module.
ValueError if the requested graph variant does not exists.
tf.errors.NotFoundError if the requested graph contains unknown ops.

variable_map Map from original variable names into tf.Variables (or lists of them).

This map translates between variable names relative to the module and the corresponding Variable objects that have been created by instantiating it in the current graph (with the applicable scoping added). Each key in the map is a variable name as created by running the module's defining module_fn in the root scope of an empty graph. Each value in the map is a Variable object, or in case of partitioned variables a list of Variable objects.

This property can be used with tf.init_from_checkpoint as assignment_map in order to restore a pre-trained checkpoint into a Module before calling Module.export().

variables Returns the list of all tf.Variables created by module instantiation.

Methods

export

View source

export(
    path, session
)

Exports the module with the variables from the session in path.

Note that it is the module definition in the ModuleSpec used to create this module that gets exported. The session is only used to provide the value of variables.

Args
path path where to export the module to.
session session where to export the variables from.

Raises
RuntimeError if there is an issue during the export.

get_attached_message

View source

get_attached_message(
    key, message_type, required=False
)

Calls ModuleSpec.get_attached_message(); see there for more.

get_input_info_dict

View source

get_input_info_dict(
    signature=None
)

Describes the inputs required by a signature.

Args
signature A string with the signature to get inputs information for. If None, the default signature is used if defined.

Returns
The result of ModuleSpec.get_input_info_dict() for the given signature, and the graph variant selected by tags when this Module was initialized.

Raises
KeyError if there is no such signature.

get_output_info_dict

View source

get_output_info_dict(
    signature=None
)

Describes the outputs provided by a signature.

Args
signature A string with the signature to get ouputs information for. If None, the default signature is used if defined.

Returns
The result of ModuleSpec.get_output_info_dict() for the given signature, and the graph variant selected by tags when this Module was initialized.

Raises
KeyError if there is no such signature.

get_signature_names

View source

get_signature_names()

Returns the module's signature names as an iterable of strings.

__call__

View source

__call__(
    inputs=None, _sentinel=None, signature=None, as_dict=None
)

Instantiates a module signature in the graph.

Example calls:

  # Use default signature with one input and default output.
  embeddings = m(["hello world", "good morning"])

  # Use "encode" signature with one input and default output.
  encodings = m(["hello world"], signature="encode")

  # Use default signature with input dict and output dict.
  dict_outputs = m({"text": [...], "lang": [...]}, as_dict=True)

The method call() allows to create the graph ops that compute a signature outputs given the inputs and using this module instance state. Each signature can be applied multiple times with different inputs and they all share the same module state.

A Module may define multiple signatures. Use signature=<name> to identify the specific signature to instantiate. If omitted or None, the default signature is used.

A signature may define various outputs. Use as_dict=True to return a dict of all outputs. If omitted or False, the output named 'default' is returned.

During this call a Module will:

  • Add ops in the current name scope to convert the inputs in tensors to feed to the signature.

  • Add ops to the UPDATE_OPS collection which depend on at least one of the provided inputs if the Module was constructed with trainable=True.

  • Add constant tensors to ASSET_FILEPATHS, even if those are not needed directly needed for the signature.

Args
inputs Inputs to the signature. A dict from input names to input tensors (incl. composite tensors, such as SparseTensor or RaggedTensor). If the signature only expects one input, one may pass a single value. If the signature has no inputs, it may be omitted.
_sentinel Used to prevent positional parameters besides inputs.
signature A string with the signature name to apply. If none, the default signature is used.
as_dict A boolean indicating whether to the return all the outputs of the signature as a dict or return only the default output.

Returns
A tensor (incl. composite tensors, such as SparseTensor or RaggedTensor) if the signature defines a default output; or a dict from strings (output names) to tensors (incl. composite tensors) if as_dict=True is used.

Raises
TypeError If there is a mismatch on arguments, inputs or outputs of the module signature.
RuntimeError If there are errors during creation of the signature graph.