def add_signature(key, inputs, outputs):
"""Adds a signature to current graph.
Args:
key: Signature key as a string.
inputs: Signature inputs as a map from string to Tensor or SparseTensor.
outputs: Signature outputs as a map from string to Tensor or SparseTensor.
(Recall that a Variable is not a Tensor, but Variable.value() is.)
Raises:
TypeError: if the arguments have the wrong types.
"""
_check_dict_maps_to_tensors_or_sparse_tensors(inputs)
_check_dict_maps_to_tensors_or_sparse_tensors(outputs)
input_info = {
input_name: tf_v1.saved_model.utils.build_tensor_info(tensor)
for input_name, tensor in inputs.items()
}
output_info = {
output_name: tf_v1.saved_model.utils.build_tensor_info(tensor)
for output_name, tensor in outputs.items()
}
signature = tf_v1.saved_model.signature_def_utils.build_signature_def(
input_info, output_info)
tf_v1.add_to_collection(_SIGNATURE_COLLECTION, (key, signature))
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 attach_bytes(key, the_bytes):
"""Adds a ModuleAttachment to the current graph.
Args:
key: A string with the unique key of the attachment.
the_bytes: A bytes object with the serialized attachment.
"""
tf_v1.add_to_collection(
_ATTACHMENT_COLLECTION_INTERNAL,
module_attachment_pb2.ModuleAttachment(key=key, value=the_bytes))
def createSavedModel(self):
model_dir = os.path.join(self.get_temp_dir(), "saved_model")
with tf.Graph().as_default():
x = tf_v1.placeholder(dtype=tf.float32, shape=[None, 3])
w = tf_v1.get_variable("weights", shape=[])
y = x * w
tf_v1.add_to_collection(_EXTRA_COLLECTION, y)
init_op = tf_v1.assign(w, 2)
with tf_v1.Session() as session:
session.run(init_op)
tf_v1.saved_model.simple_save(
session,
model_dir,
inputs={"x": x},
outputs={"y": y},
)
return model_dir