wehub-resource-sync
208 lines
8.6 KiB
Python
208 lines
8.6 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.
|
|
# ==============================================================================
|
|
"""Tests for SavedModel fingerprinting.
|
|
|
|
These tests verify that fingerprint is written correctly and that APIs for
|
|
reading it are correct.
|
|
"""
|
|
import os
|
|
import shutil
|
|
|
|
from tensorflow.core.config import flags
|
|
from tensorflow.core.protobuf import fingerprint_pb2
|
|
from tensorflow.core.protobuf import saved_model_pb2
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import test
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import tensor_spec
|
|
from tensorflow.python.lib.io import file_io
|
|
from tensorflow.python.saved_model import fingerprinting
|
|
from tensorflow.python.saved_model import fingerprinting_utils
|
|
from tensorflow.python.saved_model import save
|
|
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.trackable import autotrackable
|
|
|
|
|
|
class FingerprintingTest(test.TestCase):
|
|
|
|
def _create_saved_model(self):
|
|
root = autotrackable.AutoTrackable()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir)
|
|
self.addCleanup(shutil.rmtree, save_dir)
|
|
return save_dir
|
|
|
|
def _create_model_with_function(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(lambda x: 2. * x)
|
|
return root
|
|
|
|
def _create_model_with_input_signature(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: 2. * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
return root
|
|
|
|
def _create_model_with_data(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.x = constant_op.constant(1.0, dtype=dtypes.float32)
|
|
root.f = def_function.function(
|
|
lambda x: root.x * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
return root
|
|
|
|
def _read_fingerprint(self, filename):
|
|
fingerprint_def = fingerprint_pb2.FingerprintDef()
|
|
with file_io.FileIO(filename, "rb") as f:
|
|
fingerprint_def.ParseFromString(f.read())
|
|
return fingerprint_def
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
flags.config().saved_model_fingerprinting.reset(True)
|
|
|
|
def test_basic_module(self):
|
|
save_dir = self._create_saved_model()
|
|
files = file_io.list_directory_v2(save_dir)
|
|
|
|
self.assertLen(files, 4)
|
|
self.assertIn(constants.FINGERPRINT_FILENAME, files)
|
|
|
|
fingerprint_def = self._read_fingerprint(
|
|
file_io.join(save_dir, constants.FINGERPRINT_FILENAME))
|
|
|
|
# We cannot check this value due to non-determinism in saving.
|
|
self.assertGreater(fingerprint_def.saved_model_checksum, 0)
|
|
self.assertEqual(fingerprint_def.graph_def_program_hash,
|
|
14830488309055091319)
|
|
self.assertEqual(fingerprint_def.signature_def_hash, 12089566276354592893)
|
|
self.assertEqual(fingerprint_def.saved_object_graph_hash, 0)
|
|
# TODO(b/242348400): The checkpoint hash is non-deterministic, so we cannot
|
|
# check its value here.
|
|
self.assertGreater(fingerprint_def.checkpoint_hash, 0)
|
|
|
|
def test_model_saved_with_different_signature_options(self):
|
|
model = self._create_model_with_function()
|
|
# Save the model with signatures specified in SaveOptions.
|
|
sig_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
model,
|
|
sig_dir,
|
|
signatures=model.f.get_concrete_function(
|
|
tensor_spec.TensorSpec(None, dtypes.float32)))
|
|
# Save the model without signatures.
|
|
no_sig_dir = os.path.join(self.get_temp_dir(), "saved_model2")
|
|
save.save(model, no_sig_dir)
|
|
# Save the model with an input signature specified.
|
|
input_sig_dir = os.path.join(self.get_temp_dir(), "saved_model3")
|
|
save.save(self._create_model_with_input_signature(), input_sig_dir)
|
|
|
|
fingerprint_sig = self._read_fingerprint(
|
|
file_io.join(sig_dir, constants.FINGERPRINT_FILENAME))
|
|
fingerprint_no_sig = self._read_fingerprint(
|
|
file_io.join(no_sig_dir, constants.FINGERPRINT_FILENAME))
|
|
fingerprint_input_sig = self._read_fingerprint(
|
|
file_io.join(input_sig_dir, constants.FINGERPRINT_FILENAME))
|
|
|
|
# Check that the model saved with different options has different
|
|
# SignatureDef hashes.
|
|
self.assertNotEqual(fingerprint_sig.signature_def_hash,
|
|
fingerprint_no_sig.signature_def_hash)
|
|
# Check that the model saved with the same concrete function has the same
|
|
# regularized hashes.
|
|
self.assertEqual(fingerprint_sig.graph_def_program_hash,
|
|
fingerprint_input_sig.graph_def_program_hash)
|
|
self.assertEqual(fingerprint_sig.signature_def_hash,
|
|
fingerprint_input_sig.signature_def_hash)
|
|
|
|
def test_read_fingerprint_api(self):
|
|
save_dir = self._create_saved_model()
|
|
fingerprint = fingerprinting.read_fingerprint(save_dir)
|
|
|
|
fingerprint_def = self._read_fingerprint(
|
|
file_io.join(save_dir, constants.FINGERPRINT_FILENAME))
|
|
|
|
self.assertEqual(fingerprint, fingerprint_def)
|
|
|
|
def test_read_fingerprint_file_not_found(self):
|
|
with self.assertRaisesRegex(FileNotFoundError,
|
|
"SavedModel Fingerprint Error"):
|
|
fingerprinting.read_fingerprint("foo")
|
|
|
|
def test_write_fingerprint(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "model_and_fingerprint")
|
|
save.save_and_return_nodes(
|
|
self._create_model_with_data(), save_dir,
|
|
experimental_skip_checkpoint=True) # checkpoint data won't be loaded
|
|
|
|
fingerprint_def = fingerprinting.read_fingerprint(save_dir)
|
|
|
|
# We cannot check this value due to non-determinism in serialization.
|
|
self.assertGreater(fingerprint_def.saved_model_checksum, 0)
|
|
self.assertEqual(fingerprint_def.graph_def_program_hash,
|
|
8947653168630125217)
|
|
self.assertEqual(fingerprint_def.signature_def_hash, 15354238402988963670)
|
|
self.assertEqual(fingerprint_def.checkpoint_hash, 0)
|
|
|
|
def test_valid_singleprint(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "singleprint_model")
|
|
save.save(self._create_model_with_data(), save_dir)
|
|
fingerprint = fingerprinting.read_fingerprint(save_dir)
|
|
singleprint = fingerprint.singleprint()
|
|
|
|
# checkpoint_hash is non-deterministic and not included
|
|
self.assertRegex(singleprint,
|
|
"/".join(["8947653168630125217", # graph_def_program_hash
|
|
"15354238402988963670", # signature_def_hash
|
|
"1613952301283913051" # saved_object_graph_hash
|
|
]))
|
|
|
|
def test_invalid_singleprint(self):
|
|
fingerprint = fingerprinting.Fingerprint()
|
|
with self.assertRaisesRegex(ValueError,
|
|
"Encounted invalid fingerprint values"):
|
|
fingerprint.singleprint()
|
|
|
|
def test_valid_from_proto(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
|
|
save.save(self._create_model_with_data(), save_dir)
|
|
fingerprint_def = fingerprint_pb2.FingerprintDef().FromString(
|
|
fingerprinting_pywrap.ReadSavedModelFingerprint(save_dir))
|
|
fingerprint = fingerprinting.Fingerprint.from_proto(fingerprint_def)
|
|
self.assertEqual(fingerprint, fingerprint_def)
|
|
|
|
def test_invalid_from_proto(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
|
|
save.save(self._create_model_with_data(), save_dir)
|
|
wrong_def = saved_model_pb2.SavedModel(
|
|
saved_model_schema_version=1)
|
|
with self.assertRaisesRegex(ValueError,
|
|
"Given proto could not be deserialized as"):
|
|
fingerprinting.Fingerprint.from_proto(wrong_def)
|
|
|
|
def test_fingerprint_to_proto(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
|
|
save.save(self._create_model_with_data(), save_dir)
|
|
fingerprint = fingerprinting.read_fingerprint(save_dir)
|
|
fingerprint_def = fingerprinting_utils.to_proto(fingerprint)
|
|
self.assertEqual(fingerprint, fingerprint_def)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|