Example #1
0
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))
Example #2
0
    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
Example #3
0
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)
Example #4
0
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()
        ]