def register_module_for_export(module, export_name):
"""Register a Module to be exported under `export_name`.
DEPRECATION NOTE: This belongs to the hub.Module API and file format for TF1.
This function registers `module` to be exported by `LatestModuleExporter`
under a subdirectory named `export_name`.
Note that `export_name` must be unique for each module exported from the
current graph. It only controls the export subdirectory name and it has
no scope effects such as the `name` parameter during Module instantiation.
Args:
module: Module instance to be exported.
export_name: subdirectory name to use when performing the export.
Raises:
ValueError: if `export_name` is already taken in the current graph.
"""
for used_name, _ in tf_v1.get_collection(_EXPORT_MODULES_COLLECTION):
if used_name == export_name:
raise ValueError(
"There is already a module registered to be exported as %r" %
export_name)
tf_v1.add_to_collection(_EXPORT_MODULES_COLLECTION, (export_name, module))
def export(self,
estimator,
export_path,
checkpoint_path=None,
eval_result=None,
is_the_final_export=None):
"""Actually performs the export of registered Modules.
This method creates a timestamped directory under `export_path`
with one sub-directory (named `export_name`) per module registered
via `register_module_for_export`.
Example use:
```python
estimator = ... (Create estimator with modules registered for export)...
exporter = hub.LatestModuleExporter("tf_hub", serving_input_fn)
exporter.export(estimator, export_path, estimator.latest_checkpoint())
```
Args:
estimator: the `Estimator` from which to export modules.
export_path: A string containing a directory where to write the export
timestamped directories.
checkpoint_path: The checkpoint path to export. If `None`,
`estimator.latest_checkpoint()` is used.
eval_result: Unused.
is_the_final_export: Unused.
Returns:
The path to the created timestamped directory containing the exported
modules.
"""
if checkpoint_path is None:
checkpoint_path = estimator.latest_checkpoint()
export_dir = tf_utils.get_timestamped_export_dir(export_path)
temp_export_dir = tf_utils.get_temp_export_dir(export_dir)
session = _make_estimator_serving_session(estimator,
self._serving_input_fn,
checkpoint_path)
with session:
export_modules = tf_v1.get_collection(_EXPORT_MODULES_COLLECTION)
if export_modules:
for export_name, module in export_modules:
module_export_path = os.path.join(
temp_export_dir, tf.compat.as_bytes(export_name))
module.export(module_export_path, session)
tf_v1.gfile.Rename(temp_export_dir, export_dir)
tf_utils.garbage_collect_exports(export_path,
self._exports_to_keep)
return export_dir
else:
logging.warn(
"LatestModuleExporter found zero modules to export. "
"Use hub.register_module_for_export() if needed.")
# No export_dir has been created.
return None
def _export_signatures(meta_graph):
"""Exports signatures from current graph into a MetaGraphDef."""
named_signatures = tf_v1.get_collection(_SIGNATURE_COLLECTION)
if not named_signatures:
raise ValueError(
"No signatures present. Please call hub.add_signature(...)"
"at least once in the module_fn.")
for key, signature in named_signatures:
meta_graph.signature_def[key].CopyFrom(signature)
def _export_module_attachments(meta_graph):
"""Exports ModuleAttachments from the current tf.Graph into `meta_graph`."""
added_attachments = tf_v1.get_collection(_ATTACHMENT_COLLECTION_INTERNAL)
if not added_attachments: return # Don't touch `meta_graph`.
unique_attachments = collections.OrderedDict( # Avoid indeterminism.
(attachment.key, attachment) for attachment in added_attachments)
meta_graph.collection_def[
ATTACHMENT_COLLECTION_SAVED].bytes_list.value[:] = [
attachment.SerializeToString()
for attachment in unique_attachments.values()
]