[SavedModel Fingerprinting] Add hash 2 out of 5 for fingerprinting.

This CL adds another field (graph_def_program_hash) to the FingerprintDef protobuf.

RFC: tensorflow/community#415
PiperOrigin-RevId: 464154116

Author: Monica Song

tensorflow/cc/saved_model/fingerprinting.cc

@@ -20,6 +20,8 @@ limitations under the License.
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/op_def.pb.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/grappler/op_types.h"
 #include "tensorflow/core/lib/strings/proto_serialization.h"
 #include "tensorflow/core/platform/fingerprint.h"
@@ -31,7 +33,8 @@ namespace tensorflow::fingerprinting {
 
 namespace {
 
-// This function mutates the GraphDef, changing the names of the Function nodes.
+// This function mutates the GraphDef, changing the names and config_proto's
+// of the Function nodes.
 void CanonicalizeNodes(GraphDef* orig_graph_def) {
   for (NodeDef& node : *orig_graph_def->mutable_node()) {
     // Check if this is a function call.
@@ -42,6 +45,15 @@ void CanonicalizeNodes(GraphDef* orig_graph_def) {
       // and StatefulPartitionedCall ops.
       node.mutable_attr()->find("f")->second.mutable_func()->set_name(
           "FINGERPRINT_PASS");
+      // Erase the "config_proto" attribute which contains device-specific
+      // information.
+      node.mutable_attr()->find("config_proto")->second.mutable_s()->erase();
+    }
+    // Erase the value of string constants, which can vary based on platform.
+    if (grappler::IsConstant(node)) {
+      if (node.attr().at("dtype").type() == DT_STRING) {
+        node.mutable_attr()->find("value")->second.clear_value();
+      }
     }
   }
 }
@@ -55,8 +67,15 @@ uint64 ComputeHash(const GraphDef& graph_def) {
 }
 
 FingerprintDef CreateFingerprintDef(const MetaGraphDef& metagraph) {
+  // Create a copy of `metagraph` which will be used and mutated for fingerprint
+  // computation.
+  MetaGraphDef metagraph_copy = metagraph;
   FingerprintDef fingerprint_def;
-  fingerprint_def.set_graph_def_hash(ComputeHash(metagraph.graph_def()));
+  fingerprint_def.set_graph_def_checksum(
+      ComputeHash(metagraph_copy.graph_def()));
+  CanonicalizeGraphDef(*metagraph_copy.mutable_graph_def());
+  fingerprint_def.set_graph_def_program_hash(
+      ComputeHash(metagraph_copy.graph_def()));
   return fingerprint_def;
 }
 
@@ -67,6 +86,7 @@ void CanonicalizeGraphDef(GraphDef& graph_def) {
   // TODO(b/240173815): Complete canonicalization of the FunctionDefLibrary.
   // For now, we just clear the FunctionDefLibrary.
   graph_def.mutable_library()->Clear();
+  graph_def.mutable_versions()->Clear();
 }
 
 }  // namespace tensorflow::fingerprinting

tensorflow/cc/saved_model/fingerprinting_test.cc

@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/path.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/meta_graph.pb.h"
@@ -65,10 +66,10 @@ TEST(FingerprintingTest, TestCreateFingerprint) {
 
   TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
                           ReadSavedModel(export_dir));
-  MetaGraphDef metagraph = saved_model_pb.meta_graphs(0);
-  FingerprintDef fingerprint_def = CreateFingerprintDef(metagraph);
+  FingerprintDef fingerprint_def =
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
 
-  EXPECT_GT(fingerprint_def.graph_def_hash(), 0);
+  EXPECT_GT(fingerprint_def.graph_def_checksum(), 0);
 }
 
 // Test that canonicalization returns the same hash for two models saved by
@@ -98,5 +99,44 @@ TEST(FingerprintingTest, TestCanonicalizeGraphDeforModelSavedTwice) {
   EXPECT_EQ(hash1, hash2);
 }
 
+// Compare the fingerprints of two models saved by calling
+// `tf.saved_model.save` twice in a row in the same program.
+TEST(FingerprintingTest, TestCompareFingerprintForTwoModelSavedTwice) {
+  const std::string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
+                   "bert1", "saved_model.pb");
+
+  TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
+                          ReadSavedModel(export_dir));
+  FingerprintDef fingerprint_def =
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+
+  const std::string export_dir2 =
+      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
+                   "bert2", "saved_model.pb");
+  TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb2,
+                          ReadSavedModel(export_dir2));
+  FingerprintDef fingerprint_def2 =
+      CreateFingerprintDef(saved_model_pb2.meta_graphs(0));
+
+  EXPECT_EQ(fingerprint_def.graph_def_program_hash(),
+            fingerprint_def2.graph_def_program_hash());
+}
+
+TEST(FingerprintingTest, TestFingerprintComputationDoesNotMutateModel) {
+  const std::string export_dir =
+      io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
+                   "bert1", "saved_model.pb");
+  TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
+                          ReadSavedModel(export_dir));
+  FingerprintDef fingerprint_def =
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+  FingerprintDef fingerprint_def2 =
+      CreateFingerprintDef(saved_model_pb.meta_graphs(0));
+
+  EXPECT_EQ(fingerprint_def.graph_def_checksum(),
+            fingerprint_def2.graph_def_checksum());
+}
+
 }  // namespace
 }  // namespace tensorflow::fingerprinting

tensorflow/core/protobuf/fingerprint.proto

@@ -13,6 +13,8 @@ option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobu
 // If there are multiple MetaGraphDefs in the SavedModel, the FingerprintDef
 // corresponds to the first one.
 message FingerprintDef {
-  // Hash of the graph_def.
-  uint64 graph_def_hash = 1;
+  // Hash of the graph_def, referred to as a "checksum".
+  uint64 graph_def_checksum = 1;
+  // Hash of regularized graph_def.
+  uint64 graph_def_program_hash = 2;
 }

tensorflow/python/saved_model/fingerprinting_test.py

@@ -55,7 +55,9 @@ def test_basic_module(self):
     fingerprint_def = self._read_fingerprint(
         file_io.join(save_dir, constants.FINGERPRINT_FILENAME))
     # We cannot check the value due to non-determinism in serialization.
-    self.assertGreater(fingerprint_def.graph_def_hash, 0)
+    self.assertGreater(fingerprint_def.graph_def_checksum, 0)
+    self.assertEqual(fingerprint_def.graph_def_program_hash,
+                     16358308617800096964)
 
 
 if __name__ == "__main__":

tensorflow/python/saved_model/pywrap_saved_model_fingerprinting_test.py

@@ -31,9 +31,11 @@ def test_graphdef_basic(self):
     fingerprint_def = fingerprint_pb2.FingerprintDef()
     fingerprint_def.ParseFromString(
         fingerprinting.CreateFingerprintDef(file_content))
-    # We cannot check the value of the graph_def_hash due to non-determinism in
-    # serialization.
-    self.assertGreater(fingerprint_def.graph_def_hash, 0)
+    # We cannot check the value of the graph_def_checksum due to non-determinism
+    # in serialization.
+    self.assertGreater(fingerprint_def.graph_def_checksum, 0)
+    self.assertEqual(fingerprint_def.graph_def_program_hash,
+                     13188891313422428336)
 
 
 if __name__ == "__main__":