wehub-resource-sync
152 lines
5.2 KiB
Python
152 lines
5.2 KiB
Python
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Utilities for SavedModel fingerprinting.
|
|
|
|
This module contains utility classes and functions for working with the
|
|
SavedModel fingerprint.
|
|
"""
|
|
|
|
from absl import logging
|
|
|
|
from tensorflow.core.config import flags
|
|
from tensorflow.core.protobuf import fingerprint_pb2
|
|
from tensorflow.python.lib.io import file_io
|
|
from tensorflow.python.saved_model import fingerprinting
|
|
from tensorflow.python.saved_model.pywrap_saved_model import constants
|
|
from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as fingerprinting_pywrap
|
|
from tensorflow.python.saved_model.pywrap_saved_model import metrics
|
|
from tensorflow.python.util import compat
|
|
|
|
FingerprintException = fingerprinting_pywrap.FingerprintException
|
|
|
|
|
|
def write_fingerprint(export_dir: str) -> None:
|
|
"""Write fingerprint protobuf, if requested.
|
|
|
|
Writes a `tf.saved_model.experimental.Fingerprint` object to a
|
|
`fingerprint.pb` file in the `export_dir` using the `saved_model.pb` file
|
|
contained in `export_dir`.
|
|
|
|
Args:
|
|
export_dir: The directory in which to write the fingerprint.
|
|
"""
|
|
if flags.config().saved_model_fingerprinting.value():
|
|
fingerprint_path = file_io.join(
|
|
compat.as_str(export_dir),
|
|
compat.as_str(constants.FINGERPRINT_FILENAME))
|
|
logging.info("Writing fingerprint to %s", fingerprint_path)
|
|
|
|
try:
|
|
fingerprint_serialized = fingerprinting_pywrap.CreateFingerprintDef(
|
|
export_dir)
|
|
except FingerprintException as e:
|
|
raise ValueError(e) from None
|
|
file_io.atomic_write_string_to_file(fingerprint_path,
|
|
fingerprint_serialized)
|
|
|
|
metrics.SetWriteFingerprint(fingerprint=fingerprint_serialized)
|
|
try:
|
|
metrics.SetWritePathAndSingleprint(
|
|
path=export_dir,
|
|
singleprint=singleprint_from_fingerprint_proto(export_dir))
|
|
except metrics.MetricException:
|
|
logging.info("path_and_singleprint metric could not be set. "
|
|
"Model saving will continue.")
|
|
|
|
|
|
def singleprint_from_saved_model_proto(export_dir: str) -> str:
|
|
"""Returns the singleprint of `saved_model.pb` in `export_dir`.
|
|
|
|
Args:
|
|
export_dir: The directory that contains `saved_model.pb`.
|
|
|
|
Returns:
|
|
A string containing the singleprint of `saved_model.pb` in `export_dir`.
|
|
|
|
Raises:
|
|
ValueError: If a valid singleprint cannot be constructed from
|
|
`saved_model.pb`.
|
|
"""
|
|
try:
|
|
return fingerprinting_pywrap.SingleprintFromSM(export_dir)
|
|
except FingerprintException as e:
|
|
raise ValueError(e) from None
|
|
|
|
|
|
def singleprint_from_fingerprint_proto(export_dir: str) -> str:
|
|
"""Returns the singleprint of `fingerprint.pb` in `export_dir`.
|
|
|
|
Args:
|
|
export_dir: The directory that contains `fingerprint.pb`.
|
|
|
|
Returns:
|
|
A string containing the singleprint of `fingerprint.pb` in `export_dir`.
|
|
|
|
Raises:
|
|
ValueError: If a valid singleprint cannot be constructed from
|
|
`fingerprint.pb`.
|
|
"""
|
|
try:
|
|
return fingerprinting_pywrap.SingleprintFromFP(export_dir)
|
|
except FingerprintException as e:
|
|
raise ValueError(e) from None
|
|
|
|
|
|
def singleprint_from_saved_model(export_dir: str) -> str:
|
|
"""Returns the singleprint of the SavedModel in `export_dir`.
|
|
|
|
First tries to construct the singleprint from `fingerprint.pb`, then from
|
|
`saved_model.pb`. Attempts to write the `fingerprint.pb` if not found, but
|
|
doesn't return an error if it isn't writeable.
|
|
|
|
Args:
|
|
export_dir: The directory that contains the SavedModel.
|
|
|
|
Returns:
|
|
A string containing the singleprint of the SavedModel in `export_dir`.
|
|
|
|
Raises:
|
|
ValueError: If a valid singleprint cannot be constructed from the
|
|
SavedModel.
|
|
"""
|
|
# try generating the singleprint from `fingerprint.pb`
|
|
try:
|
|
return singleprint_from_fingerprint_proto(export_dir)
|
|
except ValueError:
|
|
pass
|
|
|
|
# try creating `fingerprint.pb`, then generating the singleprint
|
|
try:
|
|
write_fingerprint(export_dir)
|
|
return singleprint_from_fingerprint_proto(export_dir)
|
|
except ValueError:
|
|
pass
|
|
|
|
# try generating the singleprint from `saved_model.pb`
|
|
try:
|
|
return singleprint_from_saved_model_proto(export_dir)
|
|
except ValueError as e:
|
|
raise ValueError(e) from None
|
|
|
|
|
|
def to_proto(
|
|
fingerprint: fingerprinting.Fingerprint) -> fingerprint_pb2.FingerprintDef:
|
|
return fingerprint_pb2.FingerprintDef(
|
|
saved_model_checksum=fingerprint.saved_model_checksum,
|
|
graph_def_program_hash=fingerprint.graph_def_program_hash,
|
|
signature_def_hash=fingerprint.signature_def_hash,
|
|
saved_object_graph_hash=fingerprint.saved_object_graph_hash,
|
|
checkpoint_hash=fingerprint.checkpoint_hash)
|