wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/cc/saved_model/BUILD
+682
View File
@@ -0,0 +1,682 @@
#Description:
# TensorFlow SavedModel.
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load(
"//tensorflow:tensorflow.bzl",
"if_android",
"if_google",
"if_mobile",
"if_not_mobile",
"if_not_windows_or_mac",
"tf_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"if_static",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
load(
"//tensorflow/security/fuzzing:tf_fuzzing.bzl",
"tf_cc_fuzz_test",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
exports_files([
"loader.h",
"testdata/chunked_saved_model/chunked_model/saved_model.cpb",
"testdata/chunked_saved_model/chunked_model/saved_model.pbtxt",
])
cc_library(
name = "constants",
hdrs = ["constants.h"],
deps = ["@com_google_absl//absl/strings"],
)
cc_library(
name = "signature_constants",
hdrs = ["signature_constants.h"],
)
cc_library(
name = "tag_constants",
hdrs = ["tag_constants.h"],
)
# copybara:uncomment_begin(google-only)
# cc_library(
# name = "mobile_only_deps",
# visibility = ["//visibility:private"],
# deps = if_mobile(["//tensorflow/core:portable_tensorflow_lib"]),
# )
# copybara:uncomment_end
cc_library(
name = "reader",
srcs = ["reader.cc"],
hdrs = ["reader.h"],
deps = [
":constants",
":metrics",
":util",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status:statusor",
] + if_google([
"//tensorflow/tools/proto_splitter:merge",
]) + if_not_mobile([
# TODO(b/111634734): :lib and :protos_all contain dependencies that
# cannot be built on mobile platforms. Instead, include the appropriate
# tf_lib depending on the build platform.
"@com_google_absl//absl/memory:memory",
"//tensorflow/core:lib",
"//tensorflow/core/util/tensor_bundle:byteswaptensor",
]),
)
tf_cc_test(
name = "reader_test",
srcs = ["reader_test.cc"],
data = [
":saved_model_test_files",
],
linkstatic = 1,
deps = [
":constants",
":metrics",
":reader",
":tag_constants",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core/platform:resource_loader",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "loader_default_deps",
visibility = ["//visibility:private"],
deps = if_static([
"//tensorflow/core:all_kernels",
"//tensorflow/core:direct_session",
]) + if_google(if_static(["//tensorflow/core/platform/default/build_config:tensorflow_platform_specific"])),
)
cc_library(
name = "loader",
hdrs = ["loader.h"],
deps = [
":loader_lite",
] + select({
"//tensorflow:android": [],
"//tensorflow:ios": [],
"//tensorflow:tflite_converter": ["//tensorflow/core:direct_session"] + if_google([
"//tensorflow/core/platform/default/build_config:tensorflow_platform_specific",
]),
"//conditions:default": [":loader_default_deps"],
}) + if_not_mobile([
"//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
"//tensorflow/core:ops",
"//tensorflow/core:protos_all_cc",
]) + if_android([
"//tensorflow/core:portable_tensorflow_lib",
]),
)
cc_library(
name = "loader_lite",
hdrs = ["loader.h"],
deps = if_static([
":loader_lite_impl",
]) + if_not_mobile([
"//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
]),
)
cc_library(
name = "loader_lite_impl",
srcs = ["loader.cc"],
hdrs = ["loader.h"],
deps = [
":constants",
":fingerprinting",
":loader_util",
":reader",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
] + if_not_mobile([
":metrics",
":util",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/util/tensor_bundle:naming",
]),
alwayslink = 1,
)
cc_library(
name = "bundle_v2",
srcs = ["bundle_v2.cc"],
hdrs = ["bundle_v2.h"],
deps = [
":constants",
":fingerprinting",
":metrics",
":reader",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:strcat",
"//tensorflow/core/platform:tstring",
"//tensorflow/core/util/tensor_bundle",
"//tensorflow/core/util/tensor_bundle:byteswaptensor",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@jsoncpp_git//:jsoncpp",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:statusor",
"@tsl//tsl/platform:strcat",
],
)
cc_library(
name = "loader_util",
srcs = ["loader_util.cc"],
hdrs = ["loader_util.h"],
deps = [":constants"] + if_not_mobile([
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
]),
)
tf_cc_test(
name = "bundle_v2_test",
srcs = ["bundle_v2_test.cc"],
data = [
":saved_model_test_files",
],
linkstatic = 1,
deps = [
":bundle_v2",
":metrics",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:test",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@jsoncpp_git//:jsoncpp",
"@tsl//tsl/platform:statusor",
],
)
tf_cc_test(
name = "saved_model_bundle_test",
srcs = ["saved_model_bundle_test.cc"],
data = [
":saved_model_test_files",
],
linkstatic = 1,
deps = [
":constants",
":loader",
":metrics",
":reader",
":signature_constants",
":tag_constants",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
],
)
tf_cc_test(
name = "saved_model_bundle_lite_test",
srcs = ["saved_model_bundle_lite_test.cc"],
data = [
":saved_model_test_files",
],
linkstatic = 1,
deps = [
":constants",
":loader",
":signature_constants",
":tag_constants",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@com_google_absl//absl/status",
],
)
# A subset of the TF2 saved models can be generated with this tool.
py_binary(
name = "testdata/generate_saved_models",
srcs = ["testdata/generate_saved_models.py"],
data = [
":saved_model_asset_data",
":saved_model_static_hashtable_asset_data",
],
strict_deps = True,
deps = [
"//tensorflow/python/client:session",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/module",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:save_options",
"//tensorflow/python/trackable:asset",
"@absl_py//absl:app",
],
)
# copybara:uncomment_begin(google-only)
#
# py_binary(
# name = "testdata/generate_chunked_models",
# srcs = ["testdata/generate_chunked_models.py"],
# strict_deps = True,
# deps = [
# "//tensorflow/python/compat:v2_compat",
# "//tensorflow/python/eager:def_function",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/lib/io:file_io",
# "//tensorflow/python/module",
# "//tensorflow/python/saved_model:loader",
# "//tensorflow/python/saved_model:save",
# "//tensorflow/python/saved_model:save_options",
# "//tensorflow/python/util:compat",
# "//tensorflow/tools/proto_splitter:constants",
# "//tensorflow/tools/proto_splitter/python:saved_model",
# "//third_party/py/numpy",
# "@absl_py//absl:app",
# "@absl_py//absl/flags",
# ],
# )
#
# copybara:uncomment_end
# TODO(b/32673259): add a test to continuously validate these files.
filegroup(
name = "saved_model_test_files",
srcs = glob([
"testdata/AssetModule/**",
"testdata/half_plus_two_pbtxt/**",
"testdata/half_plus_two_main_op/**",
"testdata/half_plus_two/**",
"testdata/half_plus_two_v2/**",
"testdata/x_plus_y_v2_debuginfo/**",
"testdata/CyclicModule/**",
"testdata/StaticHashTableModule/**",
"testdata/VarsAndArithmeticObjectGraph/**",
"testdata/fuzz_generated/**",
"testdata/SimpleV1Model/**",
"testdata/OptimizerSlotVariableModule/**",
"testdata/chunked_saved_model/**",
]),
)
filegroup(
name = "saved_model_fingerprinting_test_files",
srcs = glob([
"testdata/bert2/**",
"testdata/bert1/**",
"testdata/half_plus_two_pbtxt/**",
]),
)
alias(
name = "saved_model_half_plus_two",
actual = ":saved_model_test_files",
)
filegroup(
name = "saved_model_asset_data",
srcs = [
"testdata/test_asset.txt",
],
)
filegroup(
name = "saved_model_static_hashtable_asset_data",
srcs = [
"testdata/static_hashtable_asset.txt",
],
)
exports_files(
glob([
"testdata/half_plus_two_pbtxt/**",
"testdata/half_plus_two_main_op/**",
"testdata/half_plus_two/**",
"testdata/half_plus_two_v2/**",
"testdata/x_plus_y_v2_debuginfo/**",
"testdata/CyclicModule/**",
"testdata/VarsAndArithmeticObjectGraph/**",
"testdata/fuzz_generated/**",
]),
)
# Linked directly into ":tensorflow_framework".
cc_library(
name = "metrics_impl",
srcs = [
"metrics.cc",
"metrics.h",
],
visibility = [
"//tensorflow:__pkg__",
"//tensorflow/python:__pkg__",
],
deps = [
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@jsoncpp_git//:jsoncpp",
] + if_not_mobile([
"//tensorflow/core:lib",
]) + if_android([
"//tensorflow/core:portable_tensorflow_lib_lite",
]),
alwayslink = True,
)
cc_library(
name = "metrics",
hdrs = ["metrics.h"],
visibility = [
"//tensorflow/cc/saved_model/image_format:__subpackages__",
"//tensorflow/python/saved_model:__subpackages__",
],
deps = if_static([
":metrics_impl",
]) + if_not_mobile([
"//tensorflow/core:lib",
]) + if_android([
"//tensorflow/core:portable_tensorflow_lib_lite",
]) + [
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
],
)
tf_cc_test(
name = "metrics_test",
size = "small",
srcs = ["metrics_test.cc"],
deps = [
":metrics",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:status_matchers",
"@com_google_absl//absl/status",
"@com_google_googletest//:gtest",
"@jsoncpp_git//:jsoncpp",
"@tsl//tsl/platform:statusor",
],
)
cc_library(
name = "util",
srcs = ["util.cc"],
hdrs = ["util.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core/framework:tensor_proto_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf:for_core_protos_cc",
] + if_not_mobile(["//tensorflow/core:lib"]) + if_android(["//tensorflow/core:portable_tensorflow_lib_lite"]),
)
cc_library(
name = "test_utils",
testonly = True,
hdrs = ["test_utils.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:test",
"//tensorflow/core/platform:protobuf",
],
)
tf_cc_test(
name = "util_test",
size = "small",
srcs = ["util_test.cc"],
deps = [
":test_utils",
":util",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/framework:tensor_shape_proto_cc",
"//tensorflow/core/framework:tensor_testutil",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@tsl//tsl/platform:status_matchers",
],
)
# Linked directly into ":tensorflow_framework".
cc_library(
name = "fingerprinting_impl",
srcs = [
"fingerprinting.cc",
"fingerprinting.h",
],
visibility = [
"//tensorflow:__pkg__",
"//tensorflow/python:__pkg__",
],
deps = [
":constants",
":fingerprinting_x_platform_utils",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/graph/regularization:simple_delete",
"//tensorflow/core/graph/regularization:util",
"//tensorflow/core/util/tensor_bundle:naming",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:types",
] + if_not_mobile([
"//tensorflow/core:lib",
]) + if_android([
"//tensorflow/core:portable_tensorflow_lib_lite",
]) + if_not_windows_or_mac([
":fingerprinting_utils",
"//tensorflow/tools/proto_splitter/cc:util",
]),
alwayslink = True,
)
cc_library(
name = "fingerprinting",
hdrs = ["fingerprinting.h"],
visibility = [
"//learning/brain/contrib/hub/server/ingestion:__subpackages__",
"//learning/brain/contrib/tpu_modeling:__subpackages__",
"//learning/gemini/deployment/disaggregation/wiz_cc:__subpackages__",
"//learning/metadata/artifactoid/cc:__subpackages__",
"//learning/tfx/pipeline/util:__subpackages__",
"//tensorflow/core/tfrt:__subpackages__",
"//tensorflow/python/saved_model:__subpackages__",
],
deps = if_static([
":fingerprinting_impl",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/status:statusor",
"//tensorflow/core:protos_all_cc",
]) + if_not_mobile([
"//tensorflow/core:lib",
]) + if_android([
"//tensorflow/core:portable_tensorflow_lib_lite",
]),
)
cc_library(
name = "fingerprinting_utils",
srcs = ["fingerprinting_utils.cc"],
hdrs = ["fingerprinting_utils.h"],
visibility = ["//visibility:private"],
deps = [
":constants",
":fingerprinting_x_platform_utils",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/util/tensor_bundle:naming",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"//tensorflow/tools/proto_splitter:merge",
"//tensorflow/tools/proto_splitter/cc:util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@riegeli//riegeli/bytes:fd_reader",
"@riegeli//riegeli/records:record_reader",
],
alwayslink = True,
)
cc_library(
name = "fingerprinting_x_platform_utils",
srcs = ["fingerprinting_x_platform_utils.cc"],
hdrs = ["fingerprinting_x_platform_utils.h"],
deps = [
"@com_google_absl//absl/numeric:int128",
"@com_google_absl//absl/strings:str_format",
"@tsl//tsl/platform:random",
],
)
tf_cc_test(
name = "fingerprinting_utils_test",
srcs = ["fingerprinting_utils_test.cc"],
data = [
"//tensorflow/tools/proto_splitter/testdata:many-field.cpb",
"//tensorflow/tools/proto_splitter/testdata:split-standard.cpb",
],
tags = [
"no_mac", # b/291933687
"no_windows", # b/291001524
],
deps = [
":fingerprinting_utils",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:path",
"//tensorflow/core/platform:protobuf",
"//tensorflow/tools/proto_splitter:chunk_proto_cc",
"//tensorflow/tools/proto_splitter/cc:util",
"//tensorflow/tools/proto_splitter/testdata:test_message_proto_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:errors",
"@tsl//tsl/platform:protobuf",
"@tsl//tsl/platform:status_matchers",
"@tsl//tsl/platform:statusor",
"@tsl//tsl/platform:test",
"@xla//xla/tsl/lib/core:status_test_util",
],
)
tf_cc_test(
name = "fingerprinting_chunked_test",
size = "small",
srcs = ["fingerprinting_chunked_test.cc"],
data = [
":saved_model_fingerprinting_test_files",
":saved_model_test_files",
],
tags = [
"no_mac", # b/291933687
"no_windows", # b/291001524
],
deps = [
":fingerprinting",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core/platform:path",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:statusor",
],
)
# copybara:uncomment_end
tf_cc_test(
name = "fingerprinting_test",
size = "small",
srcs = ["fingerprinting_test.cc"],
data = [
":saved_model_fingerprinting_test_files",
":saved_model_test_files",
],
deps = [
":fingerprinting",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/numeric:int128",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest",
],
)
tf_cc_fuzz_test(
name = "saved_model_fuzz",
srcs = ["saved_model_fuzz.cc"],
deps = [
":constants",
":loader",
":tag_constants",
"//tensorflow/core:core_cpu",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
],
)
tensorflow/cc/saved_model/README.md
+1
View File
@@ -0,0 +1 @@
<!--#include file="../../python/saved_model/README.md"-->
tensorflow/cc/saved_model/bundle_v2.cc
+234
View File
@@ -0,0 +1,234 @@
/* Copyright 2016 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/bundle_v2.h"
#include <memory>
#include <string>
#include <utility>
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/fingerprinting.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/reader.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/byte_order.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
#include "tensorflow/core/util/tensor_bundle/byte_swap_tensor.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
#include "tsl/platform/strcat.h"
namespace tensorflow {
namespace {
using strings::StrCat;
// `tensorflow::SavedModelV2Bundle::Load` API label.
constexpr char kCCLoadBundleV2Label[] = "cc_load_bundle_v2";
absl::Status ReadCheckpointObjectGraph(BundleReader* bundle_reader,
TrackableObjectGraph* object_graph) {
Tensor object_graph_tensor;
TF_RETURN_WITH_CONTEXT_IF_ERROR(
bundle_reader->Lookup(kObjectGraphProtoKey, &object_graph_tensor),
"SavedModel checkpoint does not contain object graph.");
if (object_graph_tensor.dtype() != DT_STRING ||
object_graph_tensor.dims() != 0 ||
object_graph_tensor.NumElements() != 1) {
return absl::Status(
absl::StatusCode::kFailedPrecondition,
"SavedModel checkpoint object graph was not the correct type.");
}
if (!object_graph->ParseFromString(object_graph_tensor.scalar<tstring>()())) {
return absl::Status(
absl::StatusCode::kFailedPrecondition,
"SavedModel checkpoint object graph could not be deserialized.");
}
return absl::OkStatus();
}
} // namespace
absl::Status SavedModelV2Bundle::Load(const std::string& export_dir,
SavedModelV2Bundle* const bundle) {
metrics::SavedModelReadApi(kCCLoadBundleV2Label).IncrementBy(1);
SavedModel saved_model_proto;
TF_RETURN_IF_ERROR(ReadSavedModel(export_dir, &saved_model_proto));
metrics::SavedModelReadPath().Set(export_dir);
// Load MetaGraphDef.
// In version 2 SavedModels, there is only one MetaGraphDef.
if (saved_model_proto.meta_graphs_size() != 1) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
strings::StrCat(
"SavedModelV2 should have exactly one MetaGraphDef but actually ",
"contains ", saved_model_proto.meta_graphs_size()));
}
bundle->meta_graph_def_ =
std::move(*saved_model_proto.mutable_meta_graphs(0));
// Correct the endiness of Tensor content on big-endian system
if (!port::kLittleEndian) {
TF_RETURN_IF_ERROR(
ByteSwapTensorContentInMetaGraphDef(&(bundle->meta_graph_def_)));
}
// Load GraphDebugInfo.
TF_RETURN_IF_ERROR(
ReadSavedModelDebugInfoIfPresent(export_dir, &bundle->debug_info_));
const std::string variables_dir =
io::JoinPath(export_dir, kSavedModelVariablesDirectory);
if (!Env::Default()->FileExists(variables_dir).ok()) {
LOG(INFO)
<< "No checkpoint found, assuming this is a program-only SavedModel";
} else {
// Load the variables checkpoint reader.
const std::string variables_prefix =
io::JoinPath(variables_dir, kSavedModelVariablesFilename);
bundle->variable_reader_ =
std::make_unique<BundleReader>(Env::Default(), variables_prefix);
TF_RETURN_WITH_CONTEXT_IF_ERROR(
bundle->variable_reader_->status(),
"Unable to load SavedModel variables checkpoint from ",
variables_prefix);
// Deserialize the object graph proto from the tensor bundle.
TF_RETURN_IF_ERROR(ReadCheckpointObjectGraph(
bundle->variable_reader_.get(), &bundle->trackable_object_graph_));
}
// Read the fingerprint.
auto fingerprint_proto =
saved_model::fingerprinting::ReadSavedModelFingerprint(export_dir);
if (fingerprint_proto.ok()) {
metrics::SavedModelReadFingerprint().Set(
metrics::MakeFingerprintJson(fingerprint_proto.value()));
TF_ASSIGN_OR_RETURN(
std::string path_and_singleprint,
metrics::MakeSavedModelPathAndSingleprint(
export_dir, saved_model::fingerprinting::Singleprint(
fingerprint_proto.value())));
metrics::SavedModelReadPathAndSingleprint().Set(path_and_singleprint);
}
return absl::OkStatus();
}
absl::Status SavedModelV2Bundle::VisitObjectsToRestore(
RestoreObjectsCallback callback) {
if (saved_object_graph().nodes_size() == 0 ||
trackable_object_graph().nodes_size() == 0) {
return absl::OkStatus();
}
// Start from root nodes of both the SavedObjectGraph and TrackableObjectGraph
// and descend to leaves. Note that the TrackableObjectGraph can have cycles
// (as can the SavedObjectGraph).
// This is detected and cycle edges are skipped.
const SavedObject* root_saved_object = &saved_object_graph().nodes(0);
const TrackableObjectGraph::TrackableObject* root_trackable_object =
&trackable_object_graph().nodes(0);
absl::flat_hash_set<int> trackable_node_ids;
return RecurseObjectsToRestore(root_saved_object, 0, root_trackable_object,
std::string(), &trackable_node_ids,
std::move(callback));
}
absl::Status SavedModelV2Bundle::RecurseObjectsToRestore(
const SavedObject* saved_object, int saved_object_node_id,
const TrackableObjectGraph::TrackableObject* trackable_object,
std::string object_name, absl::flat_hash_set<int>* seen_trackable_node_ids,
RestoreObjectsCallback callback) {
// Callback if any attributes or slot variables.
// Note that the root is always excluded from the search (it can never
// be a restorable object). This matches some logic on the Python side.
if (saved_object_node_id != 0 &&
(trackable_object->attributes_size() > 0 ||
trackable_object->slot_variables_size() > 0)) {
TF_RETURN_WITH_CONTEXT_IF_ERROR(
callback(saved_object_node_id, *trackable_object), "Unable to restore ",
object_name);
}
for (const auto& trackable_child_ref : trackable_object->children()) {
const auto& local_name = trackable_child_ref.local_name();
// Compute the full child name.
std::string child_name;
if (object_name.empty()) {
child_name = local_name;
} else {
child_name = strings::StrCat(object_name, ".", local_name);
}
// Descend down the trackable graph.
int trackable_child_node_id = trackable_child_ref.node_id();
if (!seen_trackable_node_ids->insert(trackable_child_node_id).second) {
// Cycle or duplicate detected - ignore this branch.
continue;
}
if (trackable_child_node_id < 0 ||
trackable_child_node_id >= trackable_object_graph().nodes_size()) {
return absl::FailedPreconditionError(
strings::StrCat("Illegal trackable child node id for ", child_name));
}
const auto* trackable_child =
&trackable_object_graph().nodes(trackable_child_node_id);
// Descend down the saved object graph.
int saved_child_node_id = -1;
const SavedObject* saved_child = nullptr;
for (const auto& saved_child_ref : saved_object->children()) {
if (saved_child_ref.local_name() == local_name) {
// Found.
saved_child_node_id = saved_child_ref.node_id();
if (saved_child_node_id >= 0 &&
saved_child_node_id < saved_object_graph().nodes_size()) {
saved_child = &saved_object_graph().nodes(saved_child_node_id);
}
break;
}
}
if (!saved_child) {
return absl::Status(
absl::StatusCode::kFailedPrecondition,
strings::StrCat("Could not find saved object to restore for ",
child_name));
}
TF_RETURN_IF_ERROR(RecurseObjectsToRestore(
saved_child, saved_child_node_id, trackable_child, child_name,
seen_trackable_node_ids, callback));
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/cc/saved_model/bundle_v2.h
+90
View File
@@ -0,0 +1,90 @@
/* Copyright 2016 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.
==============================================================================*/
// Helpers for loading the persistent representation of a SavedModelV2.
// Please note that this is depended on by code that does not make use of
// the full runtime and its dependencies should be restricted.
#ifndef TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_
#define TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_
#include <functional>
#include <memory>
#include <string>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
namespace tensorflow {
/// Represents a version 2 SavedModel that is loaded from storage (but not yet
/// loaded into an executable in-memory representation).
class SavedModelV2Bundle {
public:
using RestoreObjectsCallback = std::function<absl::Status(
int, const TrackableObjectGraph::TrackableObject&)>;
/// Loads persistent representations for a SavedModelV2 from the specified
/// export directory.
static absl::Status Load(const std::string& export_dir,
SavedModelV2Bundle* bundle);
/// MetaGraphDef from the loaded SavedModel.
MetaGraphDef& meta_graph_def() { return meta_graph_def_; }
/// SavedObjectGraph from the MetaGraphDef.
const SavedObjectGraph& saved_object_graph() {
return meta_graph_def().object_graph_def();
}
/// TrackableObjectGraph loaded from the variable_reader() checkpoint.
TrackableObjectGraph& trackable_object_graph() {
return trackable_object_graph_;
}
/// BundleReader for accessing the variables bundle.
BundleReader* variable_reader() { return variable_reader_.get(); }
/// The GraphDebugInfo (or nullptr if none).
GraphDebugInfo* debug_info() { return debug_info_.get(); }
/// Restores objects, invoking the callback with the node id in the
/// saved_object_graph() and the corresponding TrackableObject from the
/// trackable_object_graph(). The callback may use the variable_reader() but
/// must not modify the underlying saved_object_graph().
absl::Status VisitObjectsToRestore(RestoreObjectsCallback callback);
private:
absl::Status RecurseObjectsToRestore(
const SavedObject* saved_object, int saved_object_node_id,
const TrackableObjectGraph::TrackableObject* trackable_object,
std::string object_name,
absl::flat_hash_set<int>* seen_trackable_node_ids,
RestoreObjectsCallback callback);
MetaGraphDef meta_graph_def_;
TrackableObjectGraph trackable_object_graph_;
std::unique_ptr<BundleReader> variable_reader_;
std::unique_ptr<GraphDebugInfo> debug_info_;
};
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_BUNDLE_V2_H_
tensorflow/cc/saved_model/bundle_v2_test.cc
+158
View File
@@ -0,0 +1,158 @@
/* Copyright 2018 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/bundle_v2.h"
#include <algorithm>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "json/json.h"
#include "json/reader.h"
#include "json/value.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace {
constexpr char kTestData[] = "cc/saved_model/testdata";
class BundleV2Test : public ::testing::Test {
protected:
BundleV2Test() {}
void RestoreVarsAndVerify(SavedModelV2Bundle* bundle,
std::vector<std::string> expected_names) {
// Collect saved_node_id, full_name, checkpoint_key into a vector.
using RestoredVarType = std::tuple<int, std::string, std::string>;
std::vector<RestoredVarType> restored_vars;
TF_ASSERT_OK(bundle->VisitObjectsToRestore(
[&](int saved_node_id,
const TrackableObjectGraph::TrackableObject& trackable_object)
-> absl::Status {
for (const auto& attr : trackable_object.attributes()) {
if (attr.name() == "VARIABLE_VALUE") {
restored_vars.emplace_back(saved_node_id, attr.full_name(),
attr.checkpoint_key());
}
}
return absl::OkStatus();
}));
// Should be one of each var name restored.
for (const auto& expected_name : expected_names) {
EXPECT_EQ(1, std::count_if(restored_vars.begin(), restored_vars.end(),
[&](RestoredVarType t) {
return std::get<1>(t) == expected_name;
}));
}
for (const auto& restored_var : restored_vars) {
// Each restored var should match a SavedObjectGraph node with the same
// variable name.
const auto& saved_node =
bundle->saved_object_graph().nodes(std::get<0>(restored_var));
EXPECT_EQ(std::get<1>(restored_var), saved_node.variable().name());
// And should be able to load it from the tensor_bundle.
Tensor value;
TF_ASSERT_OK(
bundle->variable_reader()->Lookup(std::get<2>(restored_var), &value));
}
}
};
TEST_F(BundleV2Test, LoadsVarsAndArithmeticObjectGraph) {
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), kTestData, "VarsAndArithmeticObjectGraph");
SavedModelV2Bundle bundle;
TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
// Ensure that there are nodes in the trackable_object_graph.
EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
RestoreVarsAndVerify(&bundle, {"variable_x", "variable_y", "child_variable"});
}
TEST_F(BundleV2Test, LoadsCyclicModule) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestData, "CyclicModule");
SavedModelV2Bundle bundle;
TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
// Ensure that there are nodes in the trackable_object_graph.
EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
RestoreVarsAndVerify(&bundle, {"MyVariable"});
}
TEST_F(BundleV2Test, UpdatesMetrics) {
const std::string kCCLoadBundleV2Label = "cc_load_bundle_v2";
const int read_count = metrics::SavedModelReadCount("2").value();
const int api_count =
metrics::SavedModelReadApi(kCCLoadBundleV2Label).value();
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), kTestData, "VarsAndArithmeticObjectGraph");
SavedModelV2Bundle bundle;
TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
EXPECT_EQ(metrics::SavedModelReadCount("2").value(), read_count + 1);
EXPECT_EQ(metrics::SavedModelReadApi(kCCLoadBundleV2Label).value(),
api_count + 1);
// Check that the gauge contains the path and fingerprint.
EXPECT_EQ(metrics::SavedModelReadPath().value(), export_dir);
Json::Value fingerprint = Json::objectValue;
Json::Reader reader = Json::Reader();
reader.parse(metrics::SavedModelReadFingerprint().value(), fingerprint);
EXPECT_EQ(fingerprint["saved_model_checksum"].asUInt64(),
15788619162413586750ULL);
EXPECT_EQ(fingerprint["graph_def_program_hash"].asUInt64(),
706963557435316516ULL);
EXPECT_EQ(fingerprint["signature_def_hash"].asUInt64(),
5693392539583495303ULL);
EXPECT_EQ(fingerprint["saved_object_graph_hash"].asUInt64(),
12074714563970609759ULL);
EXPECT_EQ(fingerprint["checkpoint_hash"].asUInt64(), 10788359570789890102ULL);
TF_ASSERT_OK_AND_ASSIGN(
auto path_and_singleprint,
metrics::ParseSavedModelPathAndSingleprint(
metrics::SavedModelReadPathAndSingleprint().value()));
auto [path, singleprint] = path_and_singleprint;
EXPECT_TRUE(absl::StrContains(
path, absl::StrCat(kTestData, "/VarsAndArithmeticObjectGraph")));
EXPECT_EQ(singleprint,
"706963557435316516/" // graph_def_program_hash
"5693392539583495303/" // signature_def_hash
"12074714563970609759/" // saved_object_graph_hash
"10788359570789890102"); // checkpoint_hash
}
} // namespace
} // namespace tensorflow
tensorflow/cc/saved_model/constants.h
+82
View File
@@ -0,0 +1,82 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_
#define TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_
namespace tensorflow {
// SavedModel assets directory.
inline constexpr char kSavedModelAssetsDirectory[] = "assets";
// SavedModel assets.extra directory.
inline constexpr char kSavedModelAssetsExtraDirectory[] = "assets.extra";
// SavedModel assets key for graph collection-def.
inline constexpr char kSavedModelAssetsKey[] = "saved_model_assets";
/// SavedModel legacy init op collection key. Used in v1 SavedModels.
inline constexpr char kSavedModelLegacyInitOpKey[] = "legacy_init_op";
/// SavedModel main op collection key. Used in v1 SavedModels.
inline constexpr char kSavedModelMainOpKey[] = "saved_model_main_op";
// CollectionDef key for the SavedModel train op.
// Not exported while export_all_saved_models is experimental.
inline constexpr char kSavedModelTrainOpKey[] = "saved_model_train_op";
// Schema version for SavedModel.
inline constexpr int kSavedModelSchemaVersion = 1;
// SavedModel proto filename prefix.
inline constexpr char kSavedModelFilenamePrefix[] = "saved_model";
// SavedModel proto filename.
inline constexpr char kSavedModelFilenamePb[] = "saved_model.pb";
// SavedModel chunked proto filename.
inline constexpr char kSavedModelFilenameCpb[] = "saved_model.cpb";
// SavedModel text format proto filename.
inline constexpr char kSavedModelFilenamePbTxt[] = "saved_model.pbtxt";
// Subdirectory where debugging related files are written.
inline constexpr char kSavedModelDebugDirectory[] = "debug";
// File name for GraphDebugInfo protocol buffer which corresponds to the
// SavedModel.
inline constexpr char kSavedModelDebugInfoFilenamePb[] =
"saved_model_debug_info.pb";
// Directory in which to save the SavedModel variables.
inline constexpr char kSavedModelVariablesDirectory[] = "variables";
// SavedModel variables filename.
inline constexpr char kSavedModelVariablesFilename[] = "variables";
// SavedModel SignatureDef keys for the initialization and train ops. Used in
// V2 SavedModels.
inline constexpr char kSavedModelInitOpSignatureKey[] = "__saved_model_init_op";
inline constexpr char kSavedModelTrainOpSignatureKey[] =
"__saved_model_train_op";
// Key in the TensorBundle for the object graph proto.
inline constexpr char kObjectGraphProtoKey[] = "_CHECKPOINTABLE_OBJECT_GRAPH";
// Filename for the FingerprintDef protocol buffer.
inline constexpr char kFingerprintFilenamePb[] = "fingerprint.pb";
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_CONSTANTS_H_
tensorflow/cc/saved_model/experimental/public/BUILD
+85
View File
@@ -0,0 +1,85 @@
# Experimental C++ SavedModel Header Only APIs. See RFC
# https://github.com/tensorflow/community/pull/207
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
# This is intentionally public
default_visibility = [
"//visibility:public",
],
licenses = ["notice"],
)
cc_library(
name = "concrete_function",
hdrs = [
"concrete_function.h",
],
deps = [
":function_metadata",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/experimental/saved_model/public:concrete_function",
"//tensorflow/cc/experimental/base/public:status",
],
)
cc_library(
name = "concrete_function_list",
hdrs = [
"concrete_function_list.h",
],
deps = [
":concrete_function",
"//tensorflow/c/experimental/saved_model/public:concrete_function_list",
],
)
cc_library(
name = "function_metadata",
hdrs = [
"function_metadata.h",
],
deps = [
"//tensorflow/c/experimental/saved_model/public:function_metadata",
],
)
cc_library(
name = "saved_model_api",
hdrs = [
"saved_model_api.h",
],
deps = [
":concrete_function",
":concrete_function_list",
":signature_def_function",
"//tensorflow/c/experimental/saved_model/public:saved_model_api",
"//tensorflow/cc/experimental/base/public:runtime",
"//tensorflow/cc/experimental/base/public:status",
],
)
cc_library(
name = "signature_def_function",
hdrs = [
"signature_def_function.h",
],
deps = [
":signature_def_function_metadata",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/experimental/saved_model/public:signature_def_function",
"//tensorflow/cc/experimental/base/public:status",
],
)
cc_library(
name = "signature_def_function_metadata",
hdrs = [
"signature_def_function_metadata.h",
],
deps = [
"//tensorflow/c/experimental/saved_model/public:signature_def_function_metadata",
],
)
tensorflow/cc/saved_model/experimental/public/concrete_function.h
+61
View File
@@ -0,0 +1,61 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_
#include <vector>
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
#include "tensorflow/cc/experimental/base/public/status.h"
#include "tensorflow/cc/saved_model/experimental/public/function_metadata.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// ConcreteFunction is an executable "function" loaded from a SavedModelAPI.
class ConcreteFunction final {
public:
// TODO(bmzhao): Adding ConcreteFunction::Run in subsequent CL, since
// it depends on tensorflow::cc::Tensor and tensorflow::cc::TensorHandle
// Returns FunctionMetadata associated with this ConcreteFunction.
const FunctionMetadata* GetFunctionMetadata();
private:
friend class SavedModelAPI;
friend class ConcreteFunctionList;
// TODO(bmzhao): Consider adding a macro for wrapping/unwrapping
// when moving out of experimental.
static ConcreteFunction* wrap(TF_ConcreteFunction* p) {
return reinterpret_cast<ConcreteFunction*>(p);
}
static TF_ConcreteFunction* unwrap(ConcreteFunction* p) {
return reinterpret_cast<TF_ConcreteFunction*>(p);
}
};
inline const FunctionMetadata* ConcreteFunction::GetFunctionMetadata() {
return FunctionMetadata::wrap(TF_ConcreteFunctionGetMetadata(unwrap(this)));
}
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_H_
tensorflow/cc/saved_model/experimental/public/concrete_function_list.h
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
#include <vector>
#include "tensorflow/c/experimental/saved_model/public/concrete_function_list.h"
#include "tensorflow/cc/saved_model/experimental/public/concrete_function.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// ConcreteFunctionList helps convert an opaque pointer to an array of
// ConcreteFunction pointers to a std::vector.
class ConcreteFunctionList {
public:
// Converts this object to a std::vector<ConcreteFunction*>
std::vector<ConcreteFunction*> ToVector();
private:
friend class SavedModelAPI;
// Wraps a TF_ConcreteFunctionList. Takes ownership of list.
explicit ConcreteFunctionList(TF_ConcreteFunctionList* list) : list_(list) {}
struct TFConcreteFunctionListDeleter {
void operator()(TF_ConcreteFunctionList* p) const {
TF_DeleteConcreteFunctionList(p);
}
};
std::unique_ptr<TF_ConcreteFunctionList, TFConcreteFunctionListDeleter> list_;
};
inline std::vector<ConcreteFunction*> ConcreteFunctionList::ToVector() {
int size = TF_ConcreteFunctionListSize(list_.get());
std::vector<ConcreteFunction*> result;
result.reserve(size);
for (int i = 0; i < size; ++i) {
result.push_back(
ConcreteFunction::wrap(TF_ConcreteFunctionListGet(list_.get(), i)));
}
return result;
}
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
tensorflow/cc/saved_model/experimental/public/function_metadata.h
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_
#include <memory>
#include "tensorflow/c/experimental/saved_model/public/function_metadata.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// FunctionMetadata stores additional function information, including
// optional signaturedef feeds/fetches (for TF1-based ConcreteFunctions),
// a valid function path (for TF2-based ConcreteFunctions), and
// the types + number of inputs and outputs.
class FunctionMetadata final {
// TODO(bmzhao): Add getters here as necessary.
private:
friend class ConcreteFunction;
static FunctionMetadata* wrap(TF_FunctionMetadata* p) {
return reinterpret_cast<FunctionMetadata*>(p);
}
static TF_FunctionMetadata* unwrap(FunctionMetadata* p) {
return reinterpret_cast<TF_FunctionMetadata*>(p);
}
};
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_FUNCTION_METADATA_H_
tensorflow/cc/saved_model/experimental/public/saved_model_api.h
+155
View File
@@ -0,0 +1,155 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_
#include <memory>
#include <string>
#include <unordered_set>
#include <vector>
#include "tensorflow/c/experimental/saved_model/public/saved_model_api.h"
#include "tensorflow/cc/experimental/base/public/runtime.h"
#include "tensorflow/cc/experimental/base/public/status.h"
#include "tensorflow/cc/saved_model/experimental/public/concrete_function.h"
#include "tensorflow/cc/saved_model/experimental/public/concrete_function_list.h"
#include "tensorflow/cc/saved_model/experimental/public/signature_def_function.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// SavedModelAPI offers a way to load Tensorflow Saved Models
// (https://www.tensorflow.org/guide/saved_model) and execute saved
// tf.functions or legacy SignatureDefs in a TF2-idiomatic fashion.
// See RFC 207
// (https://github.com/tensorflow/community/blob/master/rfcs/20200218-tf-c-saved-model.md)
// TODO(bmzhao): Add an e2e example here, once ConcreteFunction::Run is added.
class SavedModelAPI {
public:
// Load a SavedModel from `dirname`.
//
// Params:
// saved_model_path - A directory filepath that the SavedModel is at.
// runtime - A runtime used to load SavedModelAPI. `runtime` must outlive the
// returned TF_SavedModel pointer.
// tags - Optional set of tags. If tags = nullptr, we expect the SavedModel
// to contain a single Metagraph (as for those exported from TF2's
// `tf.saved_model.save`). If tags != nullptr, we load the metagraph
// matching the tags:
// https://github.com/tensorflow/tensorflow/blob/428cdeda09aef81e958eeb274b83d27ad635b57b/tensorflow/core/protobuf/meta_graph.proto#L50-L56
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr.
static std::unique_ptr<SavedModelAPI> Load(
const std::string& saved_model_path, const Runtime& runtime,
Status* status, const std::unordered_set<std::string>* tags = nullptr);
// Retrieve a function from the TF2 SavedModel via function path.
//
// Params:
// function_path - A string containing the path from the root saved python
// object to a tf.function method.
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr. Otherwise, returns a
// tensorflow::cc::ConcreteFunction pointer. The lifetime of this pointer
// is bound to SavedModelAPI it was loaded from.
ConcreteFunction* GetConcreteFunction(const std::string& function_path,
Status* status);
// Retrieve a function from the TF SavedModel via a SignatureDef key.
//
// Params:
// signature_def_key - String key of SignatureDef map of a SavedModel:
// https://github.com/tensorflow/tensorflow/blob/69b08900b1e991d84bce31f3b404f5ed768f339f/tensorflow/core/protobuf/meta_graph.proto#L89
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr. Otherwise, returns a
// tensorflow::cc::ConcreteFunction pointer. The lifetime of this pointer
// is bound to SavedModelAPI it was loaded from.
SignatureDefFunction* GetSignatureDefFunction(
const std::string& function_path, Status* status);
// SavedModelAPI is movable, but not copyable.
SavedModelAPI(SavedModelAPI&&) = default;
SavedModelAPI& operator=(SavedModelAPI&&) = default;
private:
SavedModelAPI(const SavedModelAPI&) = delete;
SavedModelAPI& operator=(const SavedModelAPI&) = delete;
explicit SavedModelAPI(TF_SavedModel* model) : saved_model_(model) {}
struct TFSavedModelDeleter {
void operator()(TF_SavedModel* p) const { TF_DeleteSavedModel(p); }
};
std::unique_ptr<TF_SavedModel, TFSavedModelDeleter> saved_model_;
};
inline std::unique_ptr<SavedModelAPI> SavedModelAPI::Load(
const std::string& saved_model_path, const Runtime& runtime, Status* status,
const std::unordered_set<std::string>* tags) {
TF_SavedModel* saved_model = nullptr;
if (tags == nullptr) {
saved_model =
TF_LoadSavedModel(saved_model_path.c_str(), runtime.GetTFEContext(),
status->GetTFStatus());
} else {
std::vector<const char*> tags_vector;
tags_vector.reserve(tags->size());
for (const std::string& tag : *tags) {
tags_vector.push_back(tag.c_str());
}
saved_model = TF_LoadSavedModelWithTags(
saved_model_path.c_str(), runtime.GetTFEContext(), tags_vector.data(),
tags_vector.size(), status->GetTFStatus());
}
if (!status->ok()) {
return nullptr;
}
// We can't use std::make_unique here because of its interaction with a
// private constructor: https://abseil.io/tips/134
return std::unique_ptr<SavedModelAPI>(new SavedModelAPI(saved_model));
}
inline ConcreteFunction* SavedModelAPI::GetConcreteFunction(
const std::string& function_path, Status* status) {
TF_ConcreteFunction* function = TF_GetSavedModelConcreteFunction(
saved_model_.get(), function_path.c_str(), status->GetTFStatus());
if (!status->ok()) {
return nullptr;
}
return ConcreteFunction::wrap(function);
}
inline SignatureDefFunction* SavedModelAPI::GetSignatureDefFunction(
const std::string& function_path, Status* status) {
TF_SignatureDefFunction* function = TF_GetSavedModelSignatureDefFunction(
saved_model_.get(), function_path.c_str(), status->GetTFStatus());
if (!status->ok()) {
return nullptr;
}
return SignatureDefFunction::wrap(function);
}
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SAVED_MODEL_API_H_
tensorflow/cc/saved_model/experimental/public/signature_def_function.h
+89
View File
@@ -0,0 +1,89 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
#include <vector>
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function.h"
#include "tensorflow/cc/experimental/base/public/status.h"
#include "tensorflow/cc/saved_model/experimental/public/signature_def_function_metadata.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// SignatureDefFunctions are functions that correspond to either:
// "signatures" saved from a TF2 SavedModel APIs:
// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/python/saved_model/save.py#L830-L854
// Or the "SignatureDefMap" saved from TF1 SavedModel APIs:
// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/python/saved_model/load_v1_in_v2_test.py#L170-L174
// In both cases, a SignatureDef is serialized as a SignatureDef protobuf:
// https://github.com/tensorflow/tensorflow/blob/8ce0600f58ed84a8c84a7bbdb014d1f09e44f4c8/tensorflow/core/protobuf/meta_graph.proto#L260-L330
// and represents a computation defined by a TF subgraph.
// These Signatures were primarily designed to be interoperable with the legacy
// TF 1 Session-based C++ SavedModelBundle loading APIs:
// https://github.com/tensorflow/tensorflow/blob/26c4ee0c833e74f94d0102d8b005c41a28b44445/tensorflow/cc/saved_model/loader.h#L96-L108
// SignatureDefFunctions have different semantics from regular TF2
// ConcreteFunctions, and are mainly intended provide a serving-friendly
// transition point from the TF1 Session API.
// First, SignatureDefFunctions have different calling conventions.
// SignatureDefFunctions' inputs and outputs are constrained to **flattened
// lists of TensorHandles only**. They do not support more exotic input/output
// types (like optionals, generators, etc). Additionally, this flattening means
// they will not preserve the exact interface of the original tf.function they
// were traced from, as things like composite tensors decay into their
// internal dense tensor representation.
// Second, all inputs and outputs are "named", and these names are load bearing
// (eg: they are part of the interface of tensorflow_serving):
// https://github.com/tensorflow/serving/blob/e0d247b2e4050713194b8fad0be24a0636df7209/tensorflow_serving/apis/predict.proto#L21
// https://github.com/tensorflow/serving/blob/e0d247b2e4050713194b8fad0be24a0636df7209/tensorflow_serving/apis/predict.proto#L39
// The name of each input/output is stored in the corresponding tf::Argument in
// SignatureDefFunctionMetadata::arguments(). Users must ensure the order of
// TensorHandles passed to the function matches with the order of named
// arguments. Similarly the name of the outputs is stored in
// SignatureDefFunctionMetadata::returns().
class SignatureDefFunction final {
public:
// Returns FunctionMetadata associated with this ConcreteFunction.
const SignatureDefFunctionMetadata* GetFunctionMetadata();
private:
friend class SavedModelAPI;
friend class ConcreteFunctionList;
// TODO(bmzhao): Consider adding a macro for wrapping/unwrapping
// when moving out of experimental.
static SignatureDefFunction* wrap(TF_SignatureDefFunction* p) {
return reinterpret_cast<SignatureDefFunction*>(p);
}
static TF_SignatureDefFunction* unwrap(SignatureDefFunction* p) {
return reinterpret_cast<TF_SignatureDefFunction*>(p);
}
};
inline const SignatureDefFunctionMetadata*
SignatureDefFunction::GetFunctionMetadata() {
return SignatureDefFunctionMetadata::wrap(
TF_SignatureDefFunctionGetMetadata(unwrap(this)));
}
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
tensorflow/cc/saved_model/experimental/public/signature_def_function_metadata.h
+47
View File
@@ -0,0 +1,47 @@
/* Copyright 2020 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
#define TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
#include <memory>
#include "tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h"
namespace tensorflow {
namespace experimental {
namespace cc {
// SignatureDefFunctionMetadata stores additional information on each input
// and output's names, dtypes, and shape.
class SignatureDefFunctionMetadata final {
// TODO(bmzhao): Add getters here as necessary.
private:
friend class SignatureDefFunction;
static SignatureDefFunctionMetadata* wrap(
TF_SignatureDefFunctionMetadata* p) {
return reinterpret_cast<SignatureDefFunctionMetadata*>(p);
}
static TF_SignatureDefFunctionMetadata* unwrap(
SignatureDefFunctionMetadata* p) {
return reinterpret_cast<TF_SignatureDefFunctionMetadata*>(p);
}
};
} // namespace cc
} // namespace experimental
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_EXPERIMENTAL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
tensorflow/cc/saved_model/experimental/tests/BUILD
+28
View File
@@ -0,0 +1,28 @@
# Tests for the C++ header-only SavedModelAPI.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
tf_cc_test(
name = "saved_model_api_test",
srcs = [
"saved_model_api_test.cc",
],
data = [
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
],
deps = [
"//tensorflow/c:tf_status_headers",
"//tensorflow/cc/experimental/base/public:runtime",
"//tensorflow/cc/experimental/base/public:runtime_builder",
"//tensorflow/cc/experimental/base/public:status",
"//tensorflow/cc/saved_model/experimental/public:saved_model_api",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/strings:string_view",
],
)
tensorflow/cc/saved_model/experimental/tests/saved_model_api_test.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2020 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/experimental/public/saved_model_api.h"
#include <memory>
#include <string>
#include <unordered_set>
#include "absl/strings/string_view.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/cc/experimental/base/public/runtime.h"
#include "tensorflow/cc/experimental/base/public/runtime_builder.h"
#include "tensorflow/cc/experimental/base/public/status.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/test.h"
namespace {
using tensorflow::experimental::cc::Runtime;
using tensorflow::experimental::cc::RuntimeBuilder;
using tensorflow::experimental::cc::SavedModelAPI;
using tensorflow::experimental::cc::Status;
constexpr char kTestData[] = "cc/saved_model/testdata";
std::string SavedModelPath(absl::string_view saved_model_dir) {
return tensorflow::io::JoinPath(tensorflow::testing::TensorFlowSrcRoot(),
kTestData, saved_model_dir);
}
// This value parameterized test allows us to test both TFRT
// and non TFRT runtimes.
// https://github.com/google/googletest/blob/dcc92d0ab6c4ce022162a23566d44f673251eee4/googletest/docs/advanced.md#value-parameterized-tests
class CPPSavedModelAPITest : public ::testing::TestWithParam<bool> {};
TEST_P(CPPSavedModelAPITest, LoadsSavedModelWithTags) {
Status status;
RuntimeBuilder builder;
bool use_tfrt = GetParam();
if (use_tfrt) {
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
builder.SetUseTFRT(use_tfrt);
std::unique_ptr<Runtime> runtime = builder.Build(&status);
ASSERT_TRUE(status.ok()) << status.message();
std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
std::unordered_set<std::string> tags = {"serve"};
std::unique_ptr<SavedModelAPI> model =
SavedModelAPI::Load(model_dir, *runtime, &status, &tags);
// TODO(bmzhao): Change this to expect TF_OK when loading is implemented.
// That unblocks writing other tests that require a TF_SavedModel*,
// like loading a ConcreteFunction. This test at least checks that the
// C API builds and can be minimally run.
EXPECT_EQ(status.code(), TF_UNIMPLEMENTED);
}
TEST_P(CPPSavedModelAPITest, LoadsSavedModel) {
Status status;
RuntimeBuilder builder;
bool use_tfrt = GetParam();
if (use_tfrt) {
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
builder.SetUseTFRT(use_tfrt);
std::unique_ptr<Runtime> runtime = builder.Build(&status);
ASSERT_TRUE(status.ok()) << status.message();
std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
std::unique_ptr<SavedModelAPI> model =
SavedModelAPI::Load(model_dir, *runtime, &status);
EXPECT_EQ(status.code(), TF_OK) << status.message();
}
INSTANTIATE_TEST_SUITE_P(RuntimeAgnosticCPPSavedModelTests,
CPPSavedModelAPITest, ::testing::Bool());
} // namespace
tensorflow/cc/saved_model/fingerprinting.cc
+285
View File
@@ -0,0 +1,285 @@
/* 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/fingerprinting.h"
#include <cstdint>
#include <string>
#include "absl/container/btree_map.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/strip.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/graph/regularization/simple_delete.h"
#include "tensorflow/core/graph/regularization/util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/protobuf.h" // IWYU pragma: keep
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/util/tensor_bundle/naming.h"
// b/291933687, b/291001524
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
#include "tensorflow/cc/saved_model/fingerprinting_utils.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
#endif
// IWYU pragma: no_include "third_party/protobuf/io/coded_stream.h"
// IWYU pragma: no_include "third_party/protobuf/io/zero_copy_stream_impl_lite.h"
namespace tensorflow::saved_model::fingerprinting {
namespace {
using ::tensorflow::protobuf::Map;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::io::CodedOutputStream;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::io::StringOutputStream;
// TODO(b/290063184): remove when USM is GA
uint64_t HashCheckpointIndexFile(absl::string_view model_dir) {
std::string meta_filename = MetaFilename(io::JoinPath(
model_dir, kSavedModelVariablesDirectory, kSavedModelVariablesFilename));
std::string data;
absl::Status read_status =
ReadFileToString(Env::Default(), meta_filename, &data);
if (read_status.ok()) {
return tensorflow::Fingerprint64(data);
} else {
LOG(WARNING) << "Failed to read checkpoint file: " << read_status;
return 0;
}
}
uint64_t HashSavedModel(const SavedModel& saved_model) {
std::string saved_model_serialized;
{
// Local scope guarantees coded stream will be trimmed (ensures
// serialization determinism).
// Unfortunately the saving process itself isn't deterministic, so the
// checksum may still change since the saved_model proto may be different.
StringOutputStream stream(&saved_model_serialized);
CodedOutputStream output(&stream);
output.SetSerializationDeterministic(true);
saved_model.SerializeToCodedStream(&output);
}
return tensorflow::Fingerprint64(saved_model_serialized);
}
uint64_t RegularizeAndHashSignatureDefs(
const Map<std::string, SignatureDef>& signature_def_map) {
// Sort `signature_def_map`, which is an unordered map from string keys to
// SignatureDefs.
absl::btree_map<std::string, SignatureDef> sorted_signature_defs;
sorted_signature_defs.insert(signature_def_map.begin(),
signature_def_map.end());
uint64_t result_hash = 0;
for (const auto& item : sorted_signature_defs) {
result_hash =
FingerprintCat64(result_hash, tensorflow::Fingerprint64(item.first));
std::string signature_def_serialized;
{
StringOutputStream stream(&signature_def_serialized);
CodedOutputStream output(&stream);
output.SetSerializationDeterministic(true);
item.second.SerializeToCodedStream(&output);
}
result_hash = FingerprintCat64(
result_hash, tensorflow::Fingerprint64(signature_def_serialized));
}
return result_hash;
}
// The SavedObjectGraph contains two parts: the list of nodes and the map of
// concrete functions. Regularization treats these two parts separately.
absl::StatusOr<uint64_t> RegularizeAndHashSavedObjectGraph(
const SavedObjectGraph& object_graph_def) {
// Sort `concrete_functions`, which is an unordered map from function names to
// SavedConcreteFunction, using the suffix UID of the function name. Assumes
// that the trackable children are listed in a deterministic order during
// serialization.
absl::btree_map<int64_t, std::string> uid_to_function_names;
for (const auto& [name, concrete_function] :
object_graph_def.concrete_functions()) {
// All valid function names should end in an UID.
TF_ASSIGN_OR_RETURN(int64_t uid, graph_regularization::GetSuffixUID(name));
uid_to_function_names.insert({uid, name});
}
uint64_t result_hash = 0;
for (const auto& [uid, function_name] : uid_to_function_names) {
// Hash the function name (with the UID stripped).
result_hash = FingerprintCat64(result_hash,
tensorflow::Fingerprint64(absl::StripSuffix(
function_name, std::to_string(uid))));
// Hash the serialized concrete function.
std::string concrete_function_serialized;
{
StringOutputStream stream(&concrete_function_serialized);
CodedOutputStream output(&stream);
output.SetSerializationDeterministic(true);
object_graph_def.concrete_functions()
.at(function_name)
.SerializeToCodedStream(&output);
}
result_hash = FingerprintCat64(
result_hash, tensorflow::Fingerprint64(concrete_function_serialized));
}
// TODO(b/241294832): Complete canonicalization of `object_graph_def.nodes`.
return result_hash;
}
void SetFingerprintUUID(FingerprintDef* fingerprint_def) {
// Assign a random UUID to the fingerprint. This can happen regardless of
// whether the SavedModel was read successfully. It serves as a unique
// identifier for the model even in situations where the SavedModel is
// non-existent or unreadable.
fingerprint_def->set_uuid(CreateRandomUUID());
}
// Creates a FingerprintDef proto from a SavedModel and the checkpoint meta file
// (.index) in `export_dir`.
absl::StatusOr<FingerprintDef> CreateFingerprintDefPb(
absl::string_view export_dir, std::string pb_file) {
// Version of the code that produced the fingerprint.
// Note corresponding version definition in
// FingerprintingUtils.cc:CreateFingerprintDefCpb() and in
// Fingerprinting.cc:CreateReducedFingerprintDef()
const int kFingerprintProducer = 4;
SavedModel saved_model;
TF_RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), pb_file, &saved_model));
// Create a copy of `metagraph` which will be used and mutated for fingerprint
// computation.
FingerprintDef fingerprint_def;
MetaGraphDef* metagraph = saved_model.mutable_meta_graphs(0);
// Set fingerprint field #1.
fingerprint_def.set_saved_model_checksum(HashSavedModel(saved_model));
// Set fingerprint field #2.
graph_regularization::SimpleDelete(*metagraph->mutable_graph_def());
fingerprint_def.set_graph_def_program_hash(
graph_regularization::ComputeHash(metagraph->graph_def()));
// Set fingerprint field #3.
fingerprint_def.set_signature_def_hash(
RegularizeAndHashSignatureDefs(metagraph->signature_def()));
// Set fingerprint field #4.
TF_ASSIGN_OR_RETURN(
uint64_t object_graph_hash,
RegularizeAndHashSavedObjectGraph(metagraph->object_graph_def()));
fingerprint_def.set_saved_object_graph_hash(object_graph_hash);
// Set fingerprint field #5.
fingerprint_def.set_checkpoint_hash(HashCheckpointIndexFile(export_dir));
SetFingerprintUUID(&fingerprint_def);
// Set version of the fingerprint.
VersionDef* version = fingerprint_def.mutable_version();
version->set_producer(kFingerprintProducer);
return fingerprint_def;
}
absl::StatusOr<FingerprintDef> CreateReducedFingerprintDef() {
// Version of the code that produced the fingerprint.
// Note corresponding version definition in
// FingerprintingUtils.cc:CreateFingerprintDefCpb() and in
// Fingerprinting.cc:CreateFingerprintDefPb()
const int kFingerprintProducer = 5;
FingerprintDef fingerprint_def;
SetFingerprintUUID(&fingerprint_def);
// Set version of the fingerprint.
VersionDef* version = fingerprint_def.mutable_version();
version->set_producer(kFingerprintProducer);
return fingerprint_def;
}
} // namespace
absl::StatusOr<FingerprintDef> CreateFingerprintDef(
absl::string_view export_dir) {
std::string prefix = io::JoinPath(export_dir, kSavedModelFilenamePrefix);
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
absl::StatusOr<bool> only_contains_pb =
tools::proto_splitter::OnlyContainsPb(prefix);
// TODO: b/455882293 - Enable reporting errors if the .[c]pb files are
// important but still missing.
if (only_contains_pb.ok()) {
if (*only_contains_pb) {
return CreateFingerprintDefPb(export_dir, absl::StrCat(prefix, ".pb"));
}
return CreateFingerprintDefCpb(export_dir, absl::StrCat(prefix, ".cpb"));
}
// At this point we have neither saved_model.pb nor saved_model.cpb.
return CreateReducedFingerprintDef(); // Only sets the UUID.
#else // The following runs on Windows and Mac.
absl::StatusOr<FingerprintDef> fingerprint_def =
CreateFingerprintDefPb(export_dir, absl::StrCat(prefix, ".pb"));
if (!fingerprint_def.ok()) {
return CreateReducedFingerprintDef();
}
return fingerprint_def;
#endif
}
absl::StatusOr<FingerprintDef> ReadSavedModelFingerprint(
absl::string_view export_dir) {
const std::string fingerprint_pb_path =
io::JoinPath(export_dir, kFingerprintFilenamePb);
TF_RETURN_IF_ERROR(Env::Default()->FileExists(fingerprint_pb_path));
FingerprintDef fingerprint_proto;
absl::Status result =
ReadBinaryProto(Env::Default(), fingerprint_pb_path, &fingerprint_proto);
if (!result.ok()) return result;
return fingerprint_proto;
}
std::string Singleprint(uint64_t graph_def_program_hash,
uint64_t signature_def_hash,
uint64_t saved_object_graph_hash,
uint64_t checkpoint_hash) {
return std::to_string(graph_def_program_hash) + "/" +
std::to_string(signature_def_hash) + "/" +
std::to_string(saved_object_graph_hash) + "/" +
std::to_string(checkpoint_hash);
}
std::string Singleprint(const FingerprintDef& fingerprint) {
return Singleprint(
fingerprint.graph_def_program_hash(), fingerprint.signature_def_hash(),
fingerprint.saved_object_graph_hash(), fingerprint.checkpoint_hash());
}
absl::StatusOr<std::string> Singleprint(absl::string_view export_dir) {
TF_ASSIGN_OR_RETURN(const FingerprintDef fingerprint_def,
ReadSavedModelFingerprint(export_dir));
return Singleprint(fingerprint_def);
}
} // namespace tensorflow::saved_model::fingerprinting
tensorflow/cc/saved_model/fingerprinting.h
+47
View File
@@ -0,0 +1,47 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_
#define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_
#include <string>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
namespace tensorflow::saved_model::fingerprinting {
// Creates a FingerprintDef proto from a SavedModel (regular or chunked) and the
// checkpoint meta file (.index) in `export_dir`.
absl::StatusOr<FingerprintDef> CreateFingerprintDef(
absl::string_view export_dir);
// Loads the `fingerprint.pb` from `export_dir`, returns an error if there is
// none.
absl::StatusOr<FingerprintDef> ReadSavedModelFingerprint(
absl::string_view export_dir);
// Canonical fingerprinting ID for a SavedModel.
std::string Singleprint(uint64_t graph_def_program_hash,
uint64_t signature_def_hash,
uint64_t saved_object_graph_hash,
uint64_t checkpoint_hash);
std::string Singleprint(const FingerprintDef& fingerprint);
absl::StatusOr<std::string> Singleprint(absl::string_view export_dir);
} // namespace tensorflow::saved_model::fingerprinting
#endif // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_H_
tensorflow/cc/saved_model/fingerprinting_chunked_test.cc
+52
View File
@@ -0,0 +1,52 @@
/* Copyright 2023 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.
==============================================================================*/
#include <cstdint>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/flat_hash_set.h"
#include "tensorflow/cc/saved_model/fingerprinting.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tsl/platform/statusor.h"
namespace tensorflow::saved_model::fingerprinting {
namespace {
TEST(FingerprintingTest, TestChunkedProto) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
"chunked_saved_model/chunked_model");
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_pb,
CreateFingerprintDef(export_dir));
EXPECT_GT(fingerprint_pb.saved_model_checksum(), 0);
// We test for multiple fingerprints due to non-determinism when building with
// different compilation_mode flag options.
EXPECT_THAT(absl::flat_hash_set<uint64_t>(
{906548630859202535U, 9562420523583756263U}),
::testing::Contains(fingerprint_pb.graph_def_program_hash()));
EXPECT_EQ(fingerprint_pb.signature_def_hash(), 1043582354059066488U);
EXPECT_THAT(absl::flat_hash_set<uint64_t>(
{2766043449526180728U, 11894619660760763927U}),
::testing::Contains(fingerprint_pb.saved_object_graph_hash()));
EXPECT_EQ(fingerprint_pb.checkpoint_hash(), 0);
}
} // namespace
} // namespace tensorflow::saved_model::fingerprinting
tensorflow/cc/saved_model/fingerprinting_test.cc
+210
View File
@@ -0,0 +1,210 @@
/* 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/fingerprinting.h"
#include <string>
#include <gtest/gtest.h>
#include "absl/numeric/int128.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/numbers.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tsl/platform/statusor.h"
namespace tensorflow::saved_model::fingerprinting {
namespace {
absl::StatusOr<SavedModel> ReadSavedModel(absl::string_view file_dir) {
std::string file_path = io::JoinPath(file_dir, "saved_model.pb");
std::string serialized_saved_model;
auto status =
ReadFileToString(Env::Default(), file_path, &serialized_saved_model);
if (!status.ok()) {
return status;
}
SavedModel saved_model_pb;
saved_model_pb.ParseFromString(serialized_saved_model);
return saved_model_pb;
}
TEST(FingerprintingTest, TestCreateFingerprint) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
"VarsAndArithmeticObjectGraph");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
EXPECT_GT(fingerprint_def.saved_model_checksum(), 0);
EXPECT_EQ(fingerprint_def.graph_def_program_hash(), 10127142238652115842U);
EXPECT_EQ(fingerprint_def.signature_def_hash(), 15570736222402453744U);
EXPECT_EQ(fingerprint_def.saved_object_graph_hash(), 3678101440349108924U);
// The uuid is a random number (as string), but it should be a number > 0.
absl::uint128 uuid = 0;
EXPECT_TRUE(absl::SimpleAtoi(fingerprint_def.uuid(), &uuid))
<< "String to Uint128 conversion failed. "
<< "UUID from proto, and Uint128Max(): \n"
<< fingerprint_def.uuid() << "\n"
<< absl::Uint128Max();
EXPECT_GT(uuid, 0);
// TODO(b/242348400): The checkpoint hash is non-deterministic, so we cannot
// check its value here.
EXPECT_GT(fingerprint_def.checkpoint_hash(), 0);
}
TEST(FingerprintingTest, TestCreateFingerprintForPbtxtWorks) {
// This test ensures that we get a minimal fingerprint with an uuid, even
// when the SavedModel cannot be read.
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
"half_plus_two_pbtxt");
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
EXPECT_EQ(fingerprint_def.saved_model_checksum(), 0);
EXPECT_EQ(fingerprint_def.graph_def_program_hash(), 0);
EXPECT_EQ(fingerprint_def.signature_def_hash(), 0);
EXPECT_EQ(fingerprint_def.saved_object_graph_hash(), 0);
EXPECT_EQ(fingerprint_def.checkpoint_hash(), 0);
// The uuid is a random number (as string), but it should be a number > 0.
absl::uint128 uuid = 0;
EXPECT_TRUE(absl::SimpleAtoi(fingerprint_def.uuid(), &uuid))
<< "String to Uint128 conversion failed. "
<< "UUID from proto, and Uint128Max(): \n"
<< fingerprint_def.uuid() << "\n"
<< absl::Uint128Max();
EXPECT_GT(uuid, 0);
// version().producer()differs between WIN and non-WIN platforms.
EXPECT_GT(fingerprint_def.version().producer(), 3);
}
// 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");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
const std::string export_dir2 = io::JoinPath(
testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert2");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb2,
ReadSavedModel(export_dir2));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def2,
CreateFingerprintDef(export_dir2));
// While the saved_model serialization is deterministic, the model saving and
// proto construction is not. Therefore, we can't compare the two
// fingerprints' saved_model_checksums.
EXPECT_GT(fingerprint_def.saved_model_checksum(), 0);
EXPECT_GT(fingerprint_def2.saved_model_checksum(), 0);
EXPECT_EQ(fingerprint_def.graph_def_program_hash(),
fingerprint_def2.graph_def_program_hash());
EXPECT_EQ(fingerprint_def.signature_def_hash(),
fingerprint_def2.signature_def_hash());
EXPECT_EQ(fingerprint_def.saved_object_graph_hash(),
fingerprint_def2.saved_object_graph_hash());
EXPECT_NE(fingerprint_def.uuid(), fingerprint_def2.uuid());
}
TEST(FingerprintingTest, TestFingerprintComputationDoesNotMutateModel) {
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def2,
CreateFingerprintDef(export_dir));
EXPECT_EQ(fingerprint_def.saved_model_checksum(),
fingerprint_def2.saved_model_checksum());
}
TEST(FingerprintingTest, TestFingerprintHasVersion) {
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
EXPECT_EQ(fingerprint_def.version().producer(), 4);
}
TEST(FingerprintingTest, TestHashCheckpointForModelWithNoVariables) {
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "bert1");
TF_ASSERT_OK_AND_ASSIGN(SavedModel saved_model_pb,
ReadSavedModel(export_dir));
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_def,
CreateFingerprintDef(export_dir));
EXPECT_EQ(fingerprint_def.checkpoint_hash(), 0);
}
TEST(FingerprintingTest, TestReadValidFingerprint) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
"VarsAndArithmeticObjectGraph");
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_pb,
ReadSavedModelFingerprint(export_dir));
EXPECT_EQ(fingerprint_pb.saved_model_checksum(), 15788619162413586750u);
}
TEST(FingerprintingTest, TestReadNonexistentFingerprint) {
const std::string export_dir = io::JoinPath(
testing::TensorFlowSrcRoot(), "cc/saved_model/testdata", "AssetModule");
EXPECT_EQ(ReadSavedModelFingerprint(export_dir).status().code(),
absl::StatusCode::kNotFound);
}
TEST(FingerprintingTest, TestSingleprint) {
const std::string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
"VarsAndArithmeticObjectGraph");
const std::string const_singleprint =
"706963557435316516/5693392539583495303/12074714563970609759/"
"10788359570789890102";
TF_ASSERT_OK_AND_ASSIGN(std::string singleprint, Singleprint(export_dir));
EXPECT_EQ(singleprint, const_singleprint);
TF_ASSERT_OK_AND_ASSIGN(FingerprintDef fingerprint_pb,
ReadSavedModelFingerprint(export_dir));
EXPECT_EQ(Singleprint(fingerprint_pb), const_singleprint);
EXPECT_EQ(Singleprint(fingerprint_pb.graph_def_program_hash(),
fingerprint_pb.signature_def_hash(),
fingerprint_pb.saved_object_graph_hash(),
fingerprint_pb.checkpoint_hash()),
const_singleprint);
}
} // namespace
} // namespace tensorflow::saved_model::fingerprinting
tensorflow/cc/saved_model/fingerprinting_utils.cc
+488
View File
@@ -0,0 +1,488 @@
/* Copyright 2023 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/fingerprinting_utils.h"
#include <algorithm>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/records/record_reader.h" // from @riegeli
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/fingerprint.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/util/tensor_bundle/naming.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tensorflow/tools/proto_splitter/merge.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"
// IWYU pragma: no_include "third_party/protobuf/repeated_ptr_field.h"
// IWYU pragma: no_include "third_party/protobuf/io/coded_stream.h"
// IWYU pragma: no_include "third_party/protobuf/io/zero_copy_stream_impl_lite.h"
namespace tensorflow::saved_model::fingerprinting {
using ::tensorflow::proto_splitter::ChunkedField;
using ::tensorflow::proto_splitter::ChunkedMessage;
using ::tensorflow::proto_splitter::ChunkInfo;
using ::tensorflow::proto_splitter::ChunkMetadata;
using ::tensorflow::proto_splitter::FieldIndex;
using tools::proto_splitter::Field;
using tools::proto_splitter::FieldType;
using tools::proto_splitter::GetChunkMetadata;
using tools::proto_splitter::GetFieldTypes;
using tools::proto_splitter::GetMutableField;
using tools::proto_splitter::GetRiegeliReader;
using tools::proto_splitter::Merger;
using tools::proto_splitter::MutableFieldResult;
using tools::proto_splitter::ReadChunk;
namespace fingerprinting_utils_internal {
using ::tensorflow::protobuf::Map;
using ::tensorflow::protobuf::Message;
using ::tensorflow::protobuf::RepeatedPtrField;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::io::CodedOutputStream;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::io::StringOutputStream;
absl::StatusOr<int> fieldTagMatches(const RepeatedPtrField<FieldIndex>& a,
const RepeatedPtrField<FieldIndex>& b) {
int matches = 0;
for (int i = 0; i == matches && i < a.size() && i < b.size(); i++) {
switch (b[i].kind_case()) {
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kField:
if (a.at(i).has_field() && a.at(i).field() == b.at(i).field()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kIndex:
if (a.at(i).has_index() && a.at(i).index() == b.at(i).index()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::KindCase::kMapKey:
if (a.at(i).has_map_key()) {
const ::tensorflow::proto_splitter::FieldIndex_MapKey& key =
b.at(i).map_key();
const ::tensorflow::proto_splitter::FieldIndex_MapKey& chunked_key =
a.at(i).map_key();
switch (key.type_case()) {
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::kS:
if (chunked_key.has_s() && chunked_key.s() == key.s()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
kBoolean:
if (chunked_key.has_boolean() &&
chunked_key.boolean() == key.boolean()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
kUi32:
if (chunked_key.has_ui32() && chunked_key.ui32() == key.ui32()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
kUi64:
if (chunked_key.has_ui64() && chunked_key.ui64() == key.ui64()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
kI32:
if (chunked_key.has_i32() && chunked_key.i32() == key.i32()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
kI64:
if (chunked_key.has_i64() && chunked_key.i64() == key.i64()) {
matches += 1;
}
break;
case ::tensorflow::proto_splitter::FieldIndex::MapKey::TypeCase::
TYPE_NOT_SET:
default:
return absl::FailedPreconditionError(
"Encountered unknown field_tag.map_key type.");
}
}
break;
case FieldIndex::KindCase::KIND_NOT_SET:
default:
return absl::FailedPreconditionError(
"Encountered unknown field_tag kind.");
}
}
return matches;
}
absl::StatusOr<::tensorflow::proto_splitter::ChunkedMessage>
PruneChunkedMessage(
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
std::vector<ChunkInfo> chunks_info,
std::vector<RepeatedPtrField<FieldIndex>> target_fields_list) {
::tensorflow::proto_splitter::ChunkedMessage pruned_chunked_message;
if (chunked_message.has_chunk_index()) {
pruned_chunked_message.set_chunk_index(chunked_message.chunk_index());
}
// For each chunked_field, check if it matches any of the supplied
// target_fields, and copy over the relevant data.
for (const ChunkedField& chunked_field : chunked_message.chunked_fields()) {
for (const auto& target_fields : target_fields_list) {
TF_ASSIGN_OR_RETURN(
int matches,
fieldTagMatches(chunked_field.field_tag(), target_fields));
if (matches == chunked_field.field_tag_size()) {
// chunked_field_tags is an initial subsequence of target_fields, which
// means the chunked_field is relevant and the necessary data should be
// copied over.
auto cf = std::make_unique<proto_splitter::ChunkedField>();
cf->mutable_field_tag()->CopyFrom(chunked_field.field_tag());
TF_ASSIGN_OR_RETURN(
*cf->mutable_message(),
PruneChunkedMessage(chunked_field.message(), reader, chunks_info,
target_fields_list));
pruned_chunked_message.mutable_chunked_fields()->AddAllocated(
cf.release());
}
}
}
return pruned_chunked_message;
}
std::string SerializeProto(const Message& message) {
std::string serialized_message;
{
// local scope guarantees coded stream will be trimmed (ensures determinism)
StringOutputStream stream(&serialized_message);
CodedOutputStream output(&stream);
output.SetSerializationDeterministic(true);
message.SerializeToCodedStream(&output);
}
return serialized_message;
}
absl::StatusOr<uint64_t> HashFields(
const ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info,
const RepeatedPtrField<FieldIndex>& field_tags, Message* merged_message) {
uint64_t field_checksum = 0;
// Find chunked_fields that match the field_tags.
for (const ChunkedField& chunked_field : chunked_message.chunked_fields()) {
const RepeatedPtrField<FieldIndex> chunked_field_tags =
chunked_field.field_tag();
const ChunkedMessage& chunked_message = chunked_field.message();
// Number of sequential field_tag matches.
TF_ASSIGN_OR_RETURN(int matches,
fieldTagMatches(chunked_field_tags, field_tags));
if (chunked_message.has_chunk_index() && matches == field_tags.size()) {
// chunked_field_tags are an exact match with field_tags. Hash referenced
// chunk.
TF_ASSIGN_OR_RETURN(
std::string chunk,
ReadChunk(reader, chunks_info[chunked_message.chunk_index()]));
field_checksum = FingerprintCat64(field_checksum, Fingerprint64(chunk));
} else if (matches == field_tags.size()) {
// chunked_field_tags are an exact match, but chunked_field is further
// broken down into separate chunked_fields (no chunk_index). Hash those
// chunked_fields.
TF_ASSIGN_OR_RETURN(uint64_t hash,
HashFields(chunked_message, reader, chunks_info,
field_tags, merged_message));
field_checksum = FingerprintCat64(field_checksum, hash);
} else if (chunked_message.has_chunk_index() &&
matches == chunked_field_tags.size()) {
// chunked_field_tags are a partial match (an initial segment/subsequence
// of field_tags). Merge chunk in, attempt to locate & hash the target
// field by recursing.
TF_ASSIGN_OR_RETURN(std::vector<Field> fields,
GetFieldTypes(chunked_field_tags));
for (const auto& field : fields) {
TF_ASSIGN_OR_RETURN(MutableFieldResult mfr,
GetMutableField(merged_message, field));
merged_message =
mfr.parent->GetReflection()->MutableMessage(mfr.parent, mfr.field);
}
TF_ASSIGN_OR_RETURN(
std::string chunk,
ReadChunk(reader, chunks_info[chunked_message.chunk_index()]));
merged_message->ParseFromString(chunk);
TF_ASSIGN_OR_RETURN(uint64_t hash,
HashFields(chunked_message, reader, chunks_info,
field_tags, merged_message));
field_checksum = FingerprintCat64(field_checksum, hash);
} else if (matches == chunked_field_tags.size()) {
// chunk_field_tags are a partial match, but chunked_field is broken down.
// Merge chunked_fields in, attempt to locate & hash target field.
for (const ChunkedField& cf : chunked_message.chunked_fields()) {
TF_ASSIGN_OR_RETURN(uint64_t hash,
HashFields(cf.message(), reader, chunks_info,
field_tags, merged_message));
field_checksum = FingerprintCat64(field_checksum, hash);
}
}
}
return field_checksum;
}
inline RepeatedPtrField<FieldIndex> GraphDefFieldTags() {
// SavedModel.meta_graphs[0].graph_def
FieldIndex meta_graph_field_tag;
meta_graph_field_tag.set_field(2);
FieldIndex meta_graph_index_field_tag;
meta_graph_index_field_tag.set_index(0);
FieldIndex graph_def_field_tag;
graph_def_field_tag.set_field(2);
RepeatedPtrField<FieldIndex> graph_def_field_tags;
graph_def_field_tags.Add(FieldIndex(meta_graph_field_tag));
graph_def_field_tags.Add(FieldIndex(meta_graph_index_field_tag));
graph_def_field_tags.Add(FieldIndex(graph_def_field_tag));
return graph_def_field_tags;
}
inline RepeatedPtrField<FieldIndex> SignatureDefFieldTags() {
// SavedModel.meta_graphs[0].signature_def
FieldIndex meta_graph_field_tag;
meta_graph_field_tag.set_field(2);
FieldIndex meta_graph_index_field_tag;
meta_graph_index_field_tag.set_index(0);
FieldIndex signature_def_field_tag;
signature_def_field_tag.set_field(5);
RepeatedPtrField<FieldIndex> signature_def_field_tags;
signature_def_field_tags.Add(FieldIndex(meta_graph_field_tag));
signature_def_field_tags.Add(FieldIndex(meta_graph_index_field_tag));
signature_def_field_tags.Add(FieldIndex(signature_def_field_tag));
return signature_def_field_tags;
}
inline RepeatedPtrField<FieldIndex> SavedObjectGraphFieldTags() {
// SavedModel.meta_graphs[0].object_graph_def
FieldIndex meta_graph_field_tag;
meta_graph_field_tag.set_field(2);
FieldIndex meta_graph_index_field_tag;
meta_graph_index_field_tag.set_index(0);
FieldIndex saved_object_graph_field_tag;
saved_object_graph_field_tag.set_field(7);
RepeatedPtrField<FieldIndex> saved_object_graph_field_tags;
saved_object_graph_field_tags.Add(FieldIndex(meta_graph_field_tag));
saved_object_graph_field_tags.Add(FieldIndex(meta_graph_index_field_tag));
saved_object_graph_field_tags.Add(FieldIndex(saved_object_graph_field_tag));
return saved_object_graph_field_tags;
}
absl::StatusOr<SavedModel> PrunedSavedModel(
absl::string_view export_dir,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info, ChunkMetadata& chunk_metadata) {
SavedModel saved_model;
ChunkMetadata pruned_chunk_metadata;
pruned_chunk_metadata.mutable_chunks()->CopyFrom(chunk_metadata.chunks());
TF_ASSIGN_OR_RETURN(
*pruned_chunk_metadata.mutable_message(),
PruneChunkedMessage(chunk_metadata.message(), reader, chunks_info,
{GraphDefFieldTags(), SignatureDefFieldTags(),
SavedObjectGraphFieldTags()}));
// Read into saved_model.
TF_RETURN_IF_ERROR(
Merger::ReadPartial(io::JoinPath(export_dir, kSavedModelFilenamePrefix),
pruned_chunk_metadata, &saved_model));
return saved_model;
}
absl::StatusOr<uint64_t> HashMessage(
Message* message, const ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info,
const RepeatedPtrField<FieldIndex>& field_tags) {
uint64_t total_message_hash = Fingerprint64(SerializeProto(*message));
TF_ASSIGN_OR_RETURN(
uint64_t message_hash,
HashFields(chunked_message, reader, chunks_info, field_tags, message));
return FingerprintCat64(total_message_hash, message_hash);
}
absl::StatusOr<uint64_t> HashGraphDef(
::tensorflow::GraphDef* graph_def, const ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info) {
// TODO(adamcogdell): here we assume that graph_def (top-level) is contained
// in a single chunk, which may not be the case
return HashMessage(graph_def, chunked_message, reader, chunks_info,
GraphDefFieldTags());
}
absl::StatusOr<uint64_t> HashSignatureDef(
const Map<std::string, ::tensorflow::SignatureDef>& signature_def_map,
const ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info) {
uint64_t signature_def_hash = 0;
std::vector<std::pair<std::string, ::tensorflow::SignatureDef>>
signature_def_sorted(signature_def_map.begin(), signature_def_map.end());
std::sort(signature_def_sorted.begin(), signature_def_sorted.end(),
[](const std::pair<std::string, ::tensorflow::SignatureDef>& a,
const std::pair<std::string, ::tensorflow::SignatureDef>& b) {
return a.first < b.first;
});
for (const auto& signature_def : signature_def_sorted) {
uint64_t signature_def_pair_hash =
FingerprintCat64(Fingerprint64(signature_def.first),
Fingerprint64(SerializeProto(signature_def.second)));
signature_def_hash =
FingerprintCat64(signature_def_hash, signature_def_pair_hash);
SignatureDef signature_def_val = signature_def.second;
TF_ASSIGN_OR_RETURN(
uint64_t signature_def_entry_hash,
HashFields(chunked_message, reader, chunks_info,
SignatureDefFieldTags(), &signature_def_val));
signature_def_hash =
FingerprintCat64(signature_def_hash, signature_def_entry_hash);
}
return signature_def_hash;
}
absl::StatusOr<uint64_t> HashSavedObjectGraph(
::tensorflow::SavedObjectGraph* saved_object_graph,
const ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<ChunkInfo>& chunks_info) {
return HashMessage(saved_object_graph, chunked_message, reader, chunks_info,
SavedObjectGraphFieldTags());
}
} // namespace fingerprinting_utils_internal
using fingerprinting_utils_internal::HashFields;
using fingerprinting_utils_internal::HashGraphDef;
using fingerprinting_utils_internal::HashSavedObjectGraph;
using fingerprinting_utils_internal::HashSignatureDef;
using fingerprinting_utils_internal::PrunedSavedModel;
using fingerprinting_utils_internal::SerializeProto;
uint64_t HashCheckpointIndexFile(absl::string_view model_dir) {
std::string meta_filename = MetaFilename(io::JoinPath(
model_dir, kSavedModelVariablesDirectory, kSavedModelVariablesFilename));
std::string data;
absl::Status read_status =
ReadFileToString(Env::Default(), meta_filename, &data);
if (read_status.ok()) {
return tensorflow::Fingerprint64(data);
} else {
return 0;
}
}
absl::StatusOr<FingerprintDef> CreateFingerprintDefCpb(
absl::string_view export_dir, std::string cpb_file) {
// Version of the code that produced the fingerprint.
const int kFingerprintProducer = 2;
TF_ASSIGN_OR_RETURN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
return absl::FailedPreconditionError(
absl::StrCat("Couldn't read ChunkMetadata from chunked proto.\n",
read_metadata.status().ToString()));
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
FingerprintDef fingerprint_def;
SavedModel saved_model;
// Set the saved_model_checksum.
TF_ASSIGN_OR_RETURN(uint64_t saved_model_hash,
HashFields(chunk_metadata.message(), reader, chunks_info,
{}, &saved_model));
saved_model_hash = FingerprintCat64(
saved_model_hash, Fingerprint64(SerializeProto(saved_model)));
fingerprint_def.set_saved_model_checksum(saved_model_hash);
// Fill saved_model with only relevant chunk(s).
TF_ASSIGN_OR_RETURN(
saved_model,
PrunedSavedModel(export_dir, reader, chunks_info, chunk_metadata));
TF_ASSIGN_OR_RETURN(
uint64_t graph_def_program_hash,
HashGraphDef(saved_model.mutable_meta_graphs(0)->mutable_graph_def(),
chunk_metadata.message(), reader, chunks_info));
fingerprint_def.set_graph_def_program_hash(graph_def_program_hash);
// TODO(adamcogdell): HashSignatureDef relies on the signatue_def map being
// populated with all of its entries, which may not be the case
TF_ASSIGN_OR_RETURN(
uint64_t signature_def_hash,
HashSignatureDef(saved_model.meta_graphs(0).signature_def(),
chunk_metadata.message(), reader, chunks_info));
fingerprint_def.set_signature_def_hash(signature_def_hash);
TF_ASSIGN_OR_RETURN(
uint64_t saved_object_graph_hash,
HashSavedObjectGraph(
saved_model.mutable_meta_graphs(0)->mutable_object_graph_def(),
chunk_metadata.message(), reader, chunks_info));
fingerprint_def.set_saved_object_graph_hash(saved_object_graph_hash);
fingerprint_def.set_checkpoint_hash(HashCheckpointIndexFile(export_dir));
// Assign a random UUID to the fingerprint.
fingerprint_def.set_uuid(fingerprinting::CreateRandomUUID());
reader.Close();
// Set version of the fingerprint.
VersionDef* version = fingerprint_def.mutable_version();
version->set_producer(kFingerprintProducer);
return fingerprint_def;
}
} // namespace tensorflow::saved_model::fingerprinting
tensorflow/cc/saved_model/fingerprinting_utils.h
+137
View File
@@ -0,0 +1,137 @@
/* Copyright 2023 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_UTILS_H_
#define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_UTILS_H_
#include <cstdint>
#include <string>
#include <vector>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "riegeli/bytes/fd_reader.h" // from @riegeli
#include "riegeli/records/record_reader.h" // from @riegeli
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/protobuf.h" // IWYU pragma: keep
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
namespace tensorflow::saved_model::fingerprinting {
namespace fingerprinting_utils_internal {
using ::tensorflow::protobuf::Map;
using ::tensorflow::protobuf::Message;
using ::tensorflow::protobuf::RepeatedPtrField;
// Number of sequential FieldIndex matches of `a` in `b`. (Length of initial
// subsequence.)
// Example: `a = {4, 2}`, `b = {4, 2, 1, 3}`, `fieldTagMatches(a, b) == 2`
absl::StatusOr<int> fieldTagMatches(
const RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>& a,
const RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>& b);
// Pull out the relevant data within `chunked_message`. A `chunked_field` is
// relevant if its `field_tags` are an initial subsequence any of the
// `target_fields` in the provided `target_fields_list`.
absl::StatusOr<::tensorflow::proto_splitter::ChunkedMessage>
PruneChunkedMessage(
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
std::vector<::tensorflow::proto_splitter::ChunkInfo> chunks_info,
std::vector<RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
target_fields_list);
// Deterministically serializes the proto `message`.
std::string SerializeProto(const Message& message);
// Uses metadata contained in `chunked_message` to hash fields within the
// data accessed by the `reader` using `chunks_info`.
absl::StatusOr<uint64_t> HashFields(
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info,
const RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>&
field_tags,
Message* merged_message);
// Gets the field tags for `graph_def`.::tensorflow
inline RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>
GraphDefFieldTags();
// Gets the field tags for `signature_def`.
inline RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>
SignatureDefFieldTags();
// Gets the field tags for `saved_object_graph`.
inline RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>
SavedObjectGraphFieldTags();
// Returns a `SavedModel` containing only fields (up to those) specified by
// `GraphDefFieldTags()`, `SignatureDefFieldTags()`, and
// `SavedObjectGraphFieldTags()`.
absl::StatusOr<tensorflow::SavedModel> PrunedSavedModel(
absl::string_view export_dir,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info,
::tensorflow::proto_splitter::ChunkMetadata& chunk_metadata);
// Hashes the contents of `message` specified by `field_tags`.
absl::StatusOr<uint64_t> HashMessage(
Message* message,
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info,
const RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>&
field_tags);
// Hashes the contents of `graph_def`.
absl::StatusOr<uint64_t> HashGraphDef(
tensorflow::GraphDef* graph_def,
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info);
// Hashes the contents of `signature_def`.
absl::StatusOr<uint64_t> HashSignatureDef(
const Map<std::string, ::tensorflow::SignatureDef>& signature_def_map,
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info);
// Hashes the contents of `saved_object_graph`.
absl::StatusOr<uint64_t> HashSavedObjectGraph(
tensorflow::SavedObjectGraph* saved_object_graph,
const ::tensorflow::proto_splitter::ChunkedMessage& chunked_message,
riegeli::RecordReader<riegeli::FdReader<>>& reader,
const std::vector<::tensorflow::proto_splitter::ChunkInfo>& chunks_info);
} // namespace fingerprinting_utils_internal
// Returns the hash of the checkpoint .index file, 0 if there is none.
uint64_t HashCheckpointIndexFile(absl::string_view model_dir);
// Creates a FingerprintDef proto from a chunked SavedModel and the checkpoint
// meta file (.index) in `export_dir`.
absl::StatusOr<FingerprintDef> CreateFingerprintDefCpb(
absl::string_view export_dir, std::string cpb_file);
} // namespace tensorflow::saved_model::fingerprinting
#endif // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_UTILS_H_
tensorflow/cc/saved_model/fingerprinting_utils_test.cc
+396
View File
@@ -0,0 +1,396 @@
/* Copyright 2023 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/fingerprinting_utils.h"
#include <cstdint>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/tools/proto_splitter/cc/util.h"
#include "tensorflow/tools/proto_splitter/chunk.pb.h"
#include "tensorflow/tools/proto_splitter/testdata/test_message.pb.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/status_matchers.h"
#include "tsl/platform/statusor.h"
#include "tsl/platform/test.h"
// IWYU pragma: no_include "third_party/protobuf/io/zero_copy_stream_impl_lite.h"
// IWYU pragma: no_include "third_party/protobuf/util/message_differencer.h"
namespace tensorflow::saved_model::fingerprinting {
namespace {
using fingerprinting_utils_internal::fieldTagMatches;
using fingerprinting_utils_internal::HashFields;
using fingerprinting_utils_internal::HashGraphDef;
using fingerprinting_utils_internal::HashSavedObjectGraph;
using fingerprinting_utils_internal::HashSignatureDef;
using fingerprinting_utils_internal::PruneChunkedMessage;
using fingerprinting_utils_internal::SerializeProto;
using ::tensorflow::proto_splitter::ChunkedField;
using ::tensorflow::proto_splitter::ChunkedMessage;
using ::tensorflow::proto_splitter::ChunkInfo;
using ::tensorflow::proto_splitter::ChunkMetadata;
using ::tensorflow::proto_splitter::FieldIndex;
using ::tensorflow::proto_splitter_testdata::ManyFields;
using ::tensorflow::protobuf::Message;
using ::tensorflow::protobuf::RepeatedPtrField;
using ::tensorflow::protobuf::TextFormat;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::io::ArrayInputStream;
// NOLINTNEXTLINE: clang-tidy missing-includes false positive
using ::tensorflow::protobuf::util::MessageDifferencer;
using tools::proto_splitter::GetChunkMetadata;
using tools::proto_splitter::GetRiegeliReader;
using tsl::testing::IsOkAndHolds;
using tsl::testing::TensorFlowSrcRoot;
absl::Status ParseTextProto(absl::string_view text_proto,
Message* parsed_proto) {
TextFormat::Parser parser;
// Attempt to parse as text.
ArrayInputStream input_stream(text_proto.data(), text_proto.size());
if (parser.Parse(&input_stream, parsed_proto)) {
return absl::OkStatus();
}
parsed_proto->Clear();
return absl::InvalidArgumentError(
absl::StrCat("Could not parse text proto: ", text_proto));
}
absl::StatusOr<RepeatedPtrField<::tensorflow::proto_splitter::FieldIndex>>
ExtractFieldTags(absl::string_view chunked_field_text_proto) {
ChunkedField chunked_field;
TF_RETURN_IF_ERROR(ParseTextProto(chunked_field_text_proto, &chunked_field));
return chunked_field.field_tag();
}
TEST(FingerprintingTest, TestFieldTagMatchesInitialSubsequence) {
TF_ASSERT_OK_AND_ASSIGN(RepeatedPtrField<FieldIndex> field_tags,
ExtractFieldTags(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
RepeatedPtrField<FieldIndex> field_tags_sub;
field_tags_sub.CopyFrom(field_tags);
field_tags_sub.DeleteSubrange(2, 2);
EXPECT_THAT(fieldTagMatches(field_tags_sub, field_tags),
absl_testing::IsOkAndHolds(2));
}
TEST(FingerprintingTest, TestFieldTagMatchesNoninitialSubsequence) {
TF_ASSERT_OK_AND_ASSIGN(RepeatedPtrField<FieldIndex> field_tags,
ExtractFieldTags(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
RepeatedPtrField<FieldIndex> field_tags_sub;
field_tags_sub.CopyFrom(field_tags);
field_tags_sub.DeleteSubrange(0, 2);
EXPECT_THAT(fieldTagMatches(field_tags_sub, field_tags),
absl_testing::IsOkAndHolds(0));
}
TEST(FingerprintingTest, TestFieldTagMatchesIdenticalSubsequence) {
TF_ASSERT_OK_AND_ASSIGN(RepeatedPtrField<FieldIndex> field_tags,
ExtractFieldTags(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
RepeatedPtrField<FieldIndex> field_tags_sub;
field_tags_sub.CopyFrom(field_tags);
EXPECT_THAT(fieldTagMatches(field_tags_sub, field_tags),
absl_testing::IsOkAndHolds(4));
}
TEST(FingerprintingTest, TestFieldTagMatchesSuperSubsequence) {
TF_ASSERT_OK_AND_ASSIGN(RepeatedPtrField<FieldIndex> field_tags,
ExtractFieldTags(R"pb(
field_tag { field: 2 }
field_tag { index: 1505 }
field_tag { field: 5 }
field_tag { map_key { ui32: 123 } }
)pb"));
RepeatedPtrField<FieldIndex> field_tags_sub;
field_tags_sub.CopyFrom(field_tags);
field_tags_sub.Add()->set_field(6);
EXPECT_THAT(fieldTagMatches(field_tags_sub, field_tags),
absl_testing::IsOkAndHolds(4));
}
TEST(FingerprintingTest, TestPruneChunkedMessageSingleTarget) {
std::string cpb_file = io::JoinPath(
TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
FieldIndex field_one_field_tag;
field_one_field_tag.set_field(1);
FieldIndex repeated_field_field_tag;
repeated_field_field_tag.set_field(2);
FieldIndex repeated_field_index_field_tag;
repeated_field_index_field_tag.set_index(1);
RepeatedPtrField<FieldIndex> target_field_tags;
target_field_tags.Add(FieldIndex(field_one_field_tag));
target_field_tags.Add(FieldIndex(repeated_field_field_tag));
target_field_tags.Add(FieldIndex(repeated_field_index_field_tag));
ChunkedMessage pruned_chunked_message;
TF_ASSERT_OK_AND_ASSIGN(
pruned_chunked_message,
PruneChunkedMessage(chunk_metadata.message(), reader, chunks_info,
{target_field_tags}));
std::string expected_pruned_chunked_message_text_proto = R"pb(
chunk_index: 0
chunked_fields {
field_tag { field: 1 }
message { chunk_index: 1 }
}
)pb";
ChunkedMessage expected_pruned_chunked_message;
TF_ASSERT_OK(ParseTextProto(expected_pruned_chunked_message_text_proto,
&expected_pruned_chunked_message));
ASSERT_TRUE(MessageDifferencer::Equals(pruned_chunked_message,
expected_pruned_chunked_message));
}
TEST(FingerprintingTest, TestPruneChunkedMessageMultiTarget) {
std::string cpb_file = io::JoinPath(
TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
// ManyFields.field_one.repeated_field[1]
FieldIndex field_one_field_tag;
field_one_field_tag.set_field(1);
FieldIndex repeated_field_field_tag;
repeated_field_field_tag.set_field(2);
FieldIndex repeated_field_index_field_tag;
repeated_field_index_field_tag.set_index(1);
RepeatedPtrField<FieldIndex> target_one_field_tags;
target_one_field_tags.Add(FieldIndex(field_one_field_tag));
target_one_field_tags.Add(FieldIndex(repeated_field_field_tag));
target_one_field_tags.Add(FieldIndex(repeated_field_index_field_tag));
// ManyFields.nested_map_bool[true].string_field
FieldIndex nested_map_bool_field_tag;
nested_map_bool_field_tag.set_field(7);
FieldIndex nested_map_bool_mapkey_field_tag;
nested_map_bool_mapkey_field_tag.mutable_map_key()->set_boolean(true);
FieldIndex string_field_field_tag;
string_field_field_tag.set_field(3);
RepeatedPtrField<FieldIndex> target_two_field_tags;
target_two_field_tags.Add(FieldIndex(nested_map_bool_field_tag));
target_two_field_tags.Add(FieldIndex(nested_map_bool_mapkey_field_tag));
target_two_field_tags.Add(FieldIndex(string_field_field_tag));
ChunkedMessage pruned_chunked_message;
TF_ASSERT_OK_AND_ASSIGN(
pruned_chunked_message,
PruneChunkedMessage(chunk_metadata.message(), reader, chunks_info,
{target_one_field_tags, target_two_field_tags}));
std::string expected_pruned_chunked_message_text_proto = R"pb(
chunk_index: 0
chunked_fields {
field_tag { field: 1 }
message { chunk_index: 1 }
}
chunked_fields {
field_tag { field: 7 }
field_tag { map_key { boolean: true } }
message { chunk_index: 2 }
}
)pb";
ChunkedMessage expected_pruned_chunked_message;
TF_ASSERT_OK(ParseTextProto(expected_pruned_chunked_message_text_proto,
&expected_pruned_chunked_message));
ASSERT_TRUE(MessageDifferencer::Equals(pruned_chunked_message,
expected_pruned_chunked_message));
}
TEST(FingerprintingTest, TestPruneChunkedMessageNoTarget) {
std::string cpb_file = io::JoinPath(
TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
ChunkedMessage pruned_chunked_message;
TF_ASSERT_OK_AND_ASSIGN(
pruned_chunked_message,
PruneChunkedMessage(chunk_metadata.message(), reader, chunks_info, {}));
std::string expected_pruned_chunked_message_text_proto = R"pb(
chunk_index: 0
)pb";
ChunkedMessage expected_pruned_chunked_message;
TF_ASSERT_OK(ParseTextProto(expected_pruned_chunked_message_text_proto,
&expected_pruned_chunked_message));
ASSERT_TRUE(MessageDifferencer::Equals(pruned_chunked_message,
expected_pruned_chunked_message));
}
TEST(FingerprintingTest, TestSerializeProto) {
std::string many_fields_text_proto = R"pb(
string_field: "abc123"
)pb";
ManyFields many_fields;
TF_ASSERT_OK(ParseTextProto(many_fields_text_proto, &many_fields));
ASSERT_EQ(SerializeProto(many_fields), many_fields.SerializeAsString());
}
TEST(FingerprintingTest, TestHashFieldsV2) {
std::string cpb_file = io::JoinPath(
TensorFlowSrcRoot(), "tools/proto_splitter/testdata", "many-field.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
ManyFields many_fields;
TF_ASSERT_OK_AND_ASSIGN(uint64_t many_fields_hash,
HashFields(chunk_metadata.message(), reader,
chunks_info, {}, &many_fields));
ASSERT_EQ(many_fields_hash, 14850154939410192811U);
}
TEST(FingerprintingTest, TestHashGraphDef) {
std::string cpb_file =
io::JoinPath(TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
"split-standard.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
GraphDef graph_def;
EXPECT_THAT(
HashGraphDef(&graph_def, chunk_metadata.message(), reader, chunks_info),
absl_testing::IsOkAndHolds(16782272393894422524U));
}
TEST(FingerprintingTest, TestHashSignatureDef) {
std::string cpb_file =
io::JoinPath(TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
"split-standard.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
::tensorflow::protobuf::Map<std::string, SignatureDef> signature_def_map;
SignatureDef signature_def;
EXPECT_THAT(HashSignatureDef(signature_def_map, chunk_metadata.message(),
reader, chunks_info),
absl_testing::IsOkAndHolds(0));
}
TEST(FingerprintingTest, TestHashSavedObjectGraph) {
std::string cpb_file =
io::JoinPath(TensorFlowSrcRoot(), "tools/proto_splitter/testdata",
"split-standard.cpb");
TF_ASSERT_OK_AND_ASSIGN(auto reader, GetRiegeliReader(cpb_file));
auto read_metadata = GetChunkMetadata(reader);
if (!read_metadata.ok()) {
reader.Close();
TF_ASSERT_OK(read_metadata.status());
}
ChunkMetadata chunk_metadata = read_metadata.value();
std::vector<ChunkInfo> chunks_info = std::vector<ChunkInfo>(
chunk_metadata.chunks().begin(), chunk_metadata.chunks().end());
SavedObjectGraph saved_object_graph;
EXPECT_THAT(
HashSavedObjectGraph(&saved_object_graph, chunk_metadata.message(),
reader, chunks_info),
absl_testing::IsOkAndHolds(17454850744699451884U));
}
} // namespace
} // namespace tensorflow::saved_model::fingerprinting
tensorflow/cc/saved_model/fingerprinting_x_platform_utils.cc
+36
View File
@@ -0,0 +1,36 @@
/* Copyright 2023 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h"
#include <string>
#include "absl/numeric/int128.h"
#include "absl/strings/str_format.h"
#include "tsl/platform/random.h"
// UINT64MAX is 18'446'744'073'709'551'615 (20 digits)
// UINT128MAX is 340'282'366'920'938'463'463'374'607'431'768'211'455 (39 dgts)
// After sqrt(INT64MAX) = 4'294'967'296 (4B models), it's 50% likely to be
// duplicates in the ID space. In comparison, sqrt(UINT128MAX) = UINT64MAX,
// meaning that we can continue generating unique IDs for a lot longer time
// if the UUID is generated from two random UINT64s. This can be replaced by
// random::New128() if that becomes available.
std::string tensorflow::saved_model::fingerprinting::CreateRandomUUID() {
absl::uint128 uuid_1 = tsl::random::New64();
absl::uint128 uuid_2 = tsl::random::New64();
absl::uint128 uuid_complete = (uuid_1 << 64) | uuid_2;
return absl::StrFormat("%020d", uuid_complete);
}
tensorflow/cc/saved_model/fingerprinting_x_platform_utils.h
+28
View File
@@ -0,0 +1,28 @@
/* Copyright 2023 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_X_PLATFORM_UTILS_H_
#define TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_X_PLATFORM_UTILS_H_
#include <string>
namespace tensorflow::saved_model::fingerprinting {
// Returns a random UUID (128 bits random) as a string.
std::string CreateRandomUUID();
} // namespace tensorflow::saved_model::fingerprinting
#endif // TENSORFLOW_CC_SAVED_MODEL_FINGERPRINTING_X_PLATFORM_UTILS_H_
tensorflow/cc/saved_model/image_format/BUILD
+39
View File
@@ -0,0 +1,39 @@
# On-disk serialization format for TensorFlow models.
load(
"//tensorflow:tensorflow.bzl",
"if_not_windows_or_mac",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/cc/experimental/tfa:__subpackages__",
"//tensorflow/cc/saved_model:__subpackages__",
"//tensorflow/tools/tfg_graph_transforms:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "internal_api",
srcs = ["internal_api.cc"],
hdrs = ["internal_api.h"],
deps = [
"//tensorflow/cc/saved_model:metrics",
"//tensorflow/cc/saved_model:util",
"//tensorflow/core/platform",
"//tensorflow/core/platform:env",
"//tensorflow/core/protobuf:for_core_protos_cc",
"//tensorflow/tools/proto_splitter/cc:max_size",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:cord",
] + if_not_windows_or_mac([
"//tensorflow/tools/proto_splitter:merge",
"//tensorflow/tools/proto_splitter/cc:saved_model_splitter",
]),
)
tensorflow/cc/saved_model/image_format/README.md
+21
View File
@@ -0,0 +1,21 @@
# SavedModel Image Format
Everything related to the SavedModel Image format belongs in this directory.
If you are a TensorFlow Python user, you can try this format by setting the
`experimental_image_format` option:
```
tf.savedmodel.save(
model, path,
options=tf.saved_model.SaveOptions(experimental_image_format=True)
)
```
When this option is enabled, exported SavedModels with proto size > 2GB will
automatically save with the new format (`.cpb` instead of `.pb`).
<!-- **Compatibility** -->
The official TF APIs (TF1/TF2 python or C++ loading) have already been
integrated to handle the new format, but some downstream converters may not
have been updated.
tensorflow/cc/saved_model/image_format/internal_api.cc
+143
View File
@@ -0,0 +1,143 @@
/* Copyright 2023 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/image_format/internal_api.h"
#include <string>
#include <tuple>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/tools/proto_splitter/cc/max_size.h"
// TODO(b/291933687), TODO(b/291001524)
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
#include "tensorflow/tools/proto_splitter/cc/saved_model_splitter.h"
#include "tensorflow/tools/proto_splitter/merge.h"
#endif
#define IS_OSS false
namespace tensorflow {
namespace image_format {
absl::Status ReadSavedModel(const std::string& file_prefix,
SavedModel* saved_model_proto) {
LOG(INFO) << "Reading SavedModel from: " << file_prefix;
#if defined(PLATFORM_WINDOWS) || defined(__APPLE__)
const std::string saved_model_pb_path = absl::StrCat(file_prefix, ".pb");
TF_ASSIGN_OR_RETURN(
bool saved_model_pb_exists,
internal::FileExists(Env::Default(), saved_model_pb_path));
if (saved_model_pb_exists) {
absl::Status result =
ReadBinaryProto(Env::Default(), saved_model_pb_path, saved_model_proto);
if (result.ok()) {
metrics::SavedModelReadCount(
saved_model::GetWriteVersion(*saved_model_proto))
.IncrementBy(1);
}
return result;
}
#endif
// TODO(b/295208714): add pbtxt support to Merger::Read
const std::string saved_model_pbtxt_path =
absl::StrCat(file_prefix, ".pbtxt");
auto saved_model_pbtxt_exists =
internal::FileExists(Env::Default(), saved_model_pbtxt_path);
if (saved_model_pbtxt_exists.value_or(false)) {
absl::Status result = ReadTextProto(Env::Default(), saved_model_pbtxt_path,
saved_model_proto);
if (result.ok()) {
metrics::SavedModelReadCount(
saved_model::GetWriteVersion(*saved_model_proto))
.IncrementBy(1);
}
return result;
}
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
absl::Status result =
tools::proto_splitter::Merger::Read(file_prefix, saved_model_proto);
if (result.ok()) {
metrics::SavedModelReadCount(
saved_model::GetWriteVersion(*saved_model_proto))
.IncrementBy(1);
}
return result;
#endif
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find SavedModel .pb or .pbtxt at supplied "
"file prefix: ",
file_prefix,
". Check that "
"the directory exists and that you have the right "
"permissions for accessing it."));
}
absl::Status WriteSavedModel(SavedModel* saved_model_proto,
const std::string& file_prefix) {
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
tools::proto_splitter::SavedModelSplitter splitter(saved_model_proto);
return splitter.Write(file_prefix).status();
#else
return absl::UnimplementedError(
"WriteSavedModel not implemented for Windows or MacOS.");
#endif
}
absl::StatusOr<std::tuple<std::string, bool>> WriteSavedModelToString(
SavedModel* saved_model_proto) {
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
tools::proto_splitter::SavedModelSplitter splitter(saved_model_proto);
return splitter.WriteToString();
#else
return absl::UnimplementedError(
"WriteSavedModelToString not implemented for Windows or MacOS.");
#endif
}
#if !IS_OSS
// TODO(b/311769337): Define the function unconditionally after tf oss
// dependency is updated to protobuf v22.x.
absl::StatusOr<std::tuple<absl::Cord, bool>> WriteSavedModelToCord(
SavedModel* saved_model_proto) {
tools::proto_splitter::SavedModelSplitter splitter(saved_model_proto);
return splitter.WriteToCord();
}
#endif
absl::Status WriteSavedModel(SavedModel* saved_model_proto,
const std::string& file_prefix,
int debug_max_size) {
#if !defined(PLATFORM_WINDOWS) && !defined(__APPLE__)
tools::proto_splitter::DebugSetMaxSize(debug_max_size);
return WriteSavedModel(saved_model_proto, file_prefix);
#else
return absl::UnimplementedError(
"WriteSavedModel not implemented for Windows or MacOS.");
#endif
}
} // namespace image_format
} // namespace tensorflow
tensorflow/cc/saved_model/image_format/internal_api.h
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2023 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_IMAGE_FORMAT_INTERNAL_API_H_
#define TENSORFLOW_CC_SAVED_MODEL_IMAGE_FORMAT_INTERNAL_API_H_
#include <string>
#include <tuple>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#define IS_OSS false
namespace tensorflow {
namespace image_format {
// Reads the SavedModel proto from {file_prefix}{.pb|.cpb}.
// Returns a failure status when the SavedModel file does not exist.
absl::Status ReadSavedModel(const std::string& file_prefix,
SavedModel* saved_model_proto);
// Writes the SavedModel proto to a file or to string. If the proto is < the
// protobuf maximum size, then it will be serialized as a `.pb` proto binary.
// When larger than the maximum size, the SavedModel proto is destructively
// separated into chunks and written to
// `.cpb` (chunked proto).
//
// Write SavedModel to {file_prefix}{.pb|.cpb}.
absl::Status WriteSavedModel(SavedModel* saved_model_proto,
const std::string& file_prefix);
// Writes the SavedModel proto to std::string
// The bool field record whether it's saved as a chunked protobuf (true) or
// regular protobuf (false)
absl::StatusOr<std::tuple<std::string, bool>> WriteSavedModelToString(
SavedModel* saved_model_proto);
#if !IS_OSS
absl::StatusOr<std::tuple<absl::Cord, bool>> WriteSavedModelToCord(
SavedModel* saved_model_proto);
#endif
// See above. The `debug_max_size` argument can be used to the maximum size to
// less than 2GB for testing purposes.
absl::Status WriteSavedModel(SavedModel* saved_model_proto,
const std::string& file_prefix,
int debug_max_size);
} // namespace image_format
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_IMAGE_FORMAT_INTERNAL_API_H_
tensorflow/cc/saved_model/loader.cc
+575
View File
@@ -0,0 +1,575 @@
/* Copyright 2016 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/loader.h"
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_join.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/fingerprinting.h"
#include "tensorflow/cc/saved_model/loader_util.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/reader.h"
#include "tensorflow/cc/saved_model/util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/lib/monitoring/sampler.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/threadpool_options.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saver.pb.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/core/util/tensor_bundle/naming.h"
namespace tensorflow {
namespace {
auto* load_attempt_count = monitoring::Counter<2>::New(
"/tensorflow/cc/saved_model/load_attempt_count",
"The number of times a SavedModel was successfully loaded.", "model_path",
"status");
auto* load_latency = monitoring::Counter<1>::New(
"/tensorflow/cc/saved_model/load_latency",
"Latency in microseconds for SavedModels that were successfully loaded.",
"model_path");
auto* load_latency_by_stage = monitoring::Sampler<2>::New(
{
"/tensorflow/cc/saved_model/load_latency_by_stage", // metric name
"Distribution of wall time spent (in microseconds) in each stage "
"(restore graph from disk, run init graph op, etc) when loading the "
"model",
"model_path",
"stage",
},
// Scale of 10, power of 1.8 with bucket count 37 (~258 minutes).
monitoring::Buckets::Exponential(10, 1.8, 37));
constexpr char kLoadAttemptFail[] = "fail";
constexpr char kLoadAttemptSuccess[] = "success";
// `tensorflow::LoadSavedModel` API label.
constexpr char kCCLoadLabel[] = "cc_load";
uint64 GetLatencyMicroseconds(const uint64 start_microseconds) {
const uint64 end_microseconds = EnvTime::NowMicros();
// Avoid clock skew.
if (end_microseconds < start_microseconds) return 0;
return end_microseconds - start_microseconds;
}
// Ensure that constant tensors loaded from the saved model have valid shape.
// Also ensure that constant nodes have a value assigned to them.
// TODO(b/154763635): this is temporary and will be replaced with a better audit
static absl::Status ValidateNode(const NodeDef& node) {
const auto node_iterator = node.attr().find("value");
if (node_iterator != node.attr().end()) {
AttrValue node_value = node_iterator->second;
if (node_value.has_tensor()) {
const PartialTensorShape node_shape(node_value.tensor().tensor_shape());
if (node_shape.num_elements() < 0) {
return absl::FailedPreconditionError(absl::StrCat(
"Saved model contains node \"", node.name(), "\" (op \"", node.op(),
"\") which initializes from a tensor with ",
node_shape.num_elements(), " elements"));
}
}
} else if (node.op() == "Const") {
return absl::FailedPreconditionError(absl::StrCat(
"Saved model contains node \"", node.name(),
"\" which is a constant tensor but no value has been provided"));
}
return absl::OkStatus();
}
static absl::Status ValidateFunctionNotRecursive(const FunctionDef& function) {
const auto& function_name = function.signature().name();
for (const auto& node : function.node_def()) {
if (node.op() == function_name) {
return absl::FailedPreconditionError(absl::StrCat(
"Function ", function_name,
" is self recursive and TensorFlow does not support this scenario."));
}
}
return absl::OkStatus();
}
static absl::Status ValidateSavedTensors(const GraphDef& graph_def) {
for (const auto& node : graph_def.node()) {
TF_RETURN_IF_ERROR(ValidateNode(node));
}
if (graph_def.has_library()) {
const FunctionDefLibrary& library = graph_def.library();
for (const auto& function : library.function()) {
for (const auto& node : function.node_def()) {
TF_RETURN_IF_ERROR(ValidateNode(node));
}
// Also check that there is no recursivity in the library
TF_RETURN_IF_ERROR(ValidateFunctionNotRecursive(function));
}
}
return absl::OkStatus();
}
Tensor CreateStringTensor(const string& value) {
Tensor tensor(DT_STRING, TensorShape({}));
tensor.scalar<tstring>()() = value;
return tensor;
}
void AddAssetsTensorsToInputs(const absl::string_view export_dir,
const std::vector<AssetFileDef>& asset_file_defs,
std::vector<std::pair<string, Tensor>>* inputs) {
if (asset_file_defs.empty()) {
return;
}
for (auto& asset_file_def : asset_file_defs) {
Tensor assets_file_path_tensor = CreateStringTensor(io::JoinPath(
export_dir, kSavedModelAssetsDirectory, asset_file_def.filename()));
inputs->push_back(
{asset_file_def.tensor_info().name(), assets_file_path_tensor});
}
}
// Like Session::Run(), but uses the Make/Run/ReleaseCallable() API to avoid
// leaving behind non-GC'ed state.
//
// Detailed motivation behind this approach, from ashankar@:
//
// Each call to Session::Run() that identifies a new subgraph (based on feeds
// and fetches) creates some datastructures that live as long as the session
// (the partitioned graph, associated executors etc.).
//
// A pathological case of this would be if say the initialization op
// (main_op/legacy_init_op) involves the use of a large constant. Then we
// allocate memory for that large constant that will just stick around till the
// session dies. With this Callable mechanism, that memory will be released
// right after ReleaseCallable returns.
//
// However, the resource manager state remains.
absl::Status RunOnce(const RunOptions& run_options,
const std::vector<std::pair<string, Tensor>>& inputs,
const std::vector<string>& output_tensor_names,
const std::vector<string>& target_node_names,
std::vector<Tensor>* outputs, RunMetadata* run_metadata,
Session* session) {
CallableOptions callable_options;
std::vector<Tensor> feed_tensors;
*callable_options.mutable_run_options() = run_options;
for (const auto& input : inputs) {
const string& name = input.first;
const Tensor& tensor = input.second;
callable_options.add_feed(name);
feed_tensors.push_back(tensor);
}
for (const string& output_tensor_name : output_tensor_names) {
callable_options.add_fetch(output_tensor_name);
}
for (const string& target_node_name : target_node_names) {
callable_options.add_target(target_node_name);
}
Session::CallableHandle callable_handle;
TF_RETURN_IF_ERROR(session->MakeCallable(callable_options, &callable_handle));
const absl::Status run_status = session->RunCallable(
callable_handle, feed_tensors, outputs, run_metadata);
// Be sure to call ReleaseCallable() regardless of the outcome of
// RunCallable().
session->ReleaseCallable(callable_handle).IgnoreError();
return run_status;
}
// RunInitOp will return OK if the initialization op was run successfully.
// An empty init_op_name indicates that there are no init ops to run.
absl::Status RunInitOp(const RunOptions& run_options, const string& export_dir,
const MetaGraphDef& meta_graph_def,
const std::vector<AssetFileDef>& asset_file_defs,
Session* session, const string& init_op_name) {
if (!init_op_name.empty()) {
LOG(INFO) << "Running initialization op on SavedModel bundle at path: "
<< export_dir;
std::vector<std::pair<string, Tensor>> inputs;
AddAssetsTensorsToInputs(export_dir, asset_file_defs, &inputs);
RunMetadata run_metadata;
return RunOnce(run_options, inputs, {}, {init_op_name},
nullptr /* outputs */, &run_metadata, session);
}
return absl::OkStatus();
}
absl::Status RunRestore(const RunOptions& run_options, const string& export_dir,
const absl::string_view restore_op_name,
const absl::string_view variable_filename_const_op_name,
const std::vector<AssetFileDef>& asset_file_defs,
Session* session) {
LOG(INFO) << "Restoring SavedModel bundle.";
// Find path to variables to be restored in export directory.
const string variables_directory =
io::JoinPath(export_dir, kSavedModelVariablesDirectory);
// Check for saver checkpoints in v2 format. Models exported in the checkpoint
// v2 format will have a variables.index file. The corresponding
// variables are stored in the variables.data-?????-of-????? files.
const string variables_index_path = io::JoinPath(
variables_directory, MetaFilename(kSavedModelVariablesFilename));
TF_ASSIGN_OR_RETURN(
bool variables_index_exists,
internal::FileExists(Env::Default(), variables_index_path));
if (!variables_index_exists) {
LOG(INFO) << "The specified SavedModel has no variables; no checkpoints "
"were restored. File does not exist: "
<< variables_index_path;
return absl::OkStatus();
}
const string variables_path =
io::JoinPath(variables_directory, kSavedModelVariablesFilename);
// Add variables to the graph.
Tensor variables_path_tensor(DT_STRING, TensorShape({}));
variables_path_tensor.scalar<tstring>()() = variables_path;
std::vector<std::pair<string, Tensor>> inputs = {
{string(variable_filename_const_op_name), variables_path_tensor}};
AddAssetsTensorsToInputs(export_dir, asset_file_defs, &inputs);
RunMetadata run_metadata;
return RunOnce(run_options, inputs, {}, {string(restore_op_name)},
nullptr /* outputs */, &run_metadata, session);
}
} // namespace
SavedModelBundleInterface::~SavedModelBundleInterface() = default;
absl::Status LoadMetagraphIntoSession(const SessionOptions& session_options,
const MetaGraphDef& meta_graph,
std::unique_ptr<Session>* session) {
Session* session_p = nullptr;
TF_RETURN_IF_ERROR(NewSession(session_options, &session_p));
session->reset(session_p);
TF_RETURN_IF_ERROR(ValidateSavedTensors(meta_graph.graph_def()));
return (*session)->Create(meta_graph.graph_def());
}
absl::Status LoadGraphDefIntoSession(const SessionOptions& session_options,
GraphDef graph_def,
std::unique_ptr<Session>* session) {
Session* session_p = nullptr;
TF_RETURN_IF_ERROR(NewSession(session_options, &session_p));
session->reset(session_p);
TF_RETURN_IF_ERROR(ValidateSavedTensors(graph_def));
return (*session)->Create(std::move(graph_def));
}
absl::Status LoadSavedModelInternal(const SessionOptions& session_options,
const RunOptions& run_options,
const string& export_dir,
const std::unordered_set<string>& tags,
SavedModelBundle* const bundle) {
TF_RETURN_IF_ERROR(ReadMetaGraphDefFromSavedModel(export_dir, tags,
&bundle->meta_graph_def));
TF_RETURN_IF_ERROR(
ReadSavedModelDebugInfoIfPresent(export_dir, &bundle->debug_info));
TF_RETURN_IF_ERROR(LoadMetagraphIntoSession(
session_options, bundle->meta_graph_def, &bundle->session));
TF_RETURN_IF_ERROR(RestoreSession(run_options, bundle->meta_graph_def,
export_dir, &bundle->session));
return absl::OkStatus();
}
namespace {
// Session wrapper that prevents calls to Session::Create(), Session::Extend(),
// and the deprecated partial-run methods.
//
// Limiting the available methods on a returned Session gives us the option
// to replace the Session with a cut-down implementation, without breaking any
// users.
class LiteSessionWrapper : public Session {
public:
explicit LiteSessionWrapper(std::unique_ptr<Session> wrapped)
: wrapped_(std::move(wrapped)) {}
absl::Status Create(const GraphDef& graph) override {
return absl::UnimplementedError("Session::Create()");
}
absl::Status Create(GraphDef&& graph) override {
return absl::UnimplementedError("Session::Create()");
}
absl::Status Extend(const GraphDef& graph) override {
return absl::UnimplementedError("Session::Extend()");
}
absl::Status Extend(GraphDef&& graph) override {
return absl::UnimplementedError("Session::Extend()");
}
absl::Status Run(const std::vector<std::pair<string, Tensor>>& inputs,
const std::vector<string>& output_tensor_names,
const std::vector<string>& target_node_names,
std::vector<Tensor>* outputs) override {
return wrapped_->Run(inputs, output_tensor_names, target_node_names,
outputs);
}
absl::Status Create(const RunOptions& run_options,
const GraphDef& graph) override {
return absl::UnimplementedError("Session::Create()");
}
absl::Status Extend(const RunOptions& run_options,
const GraphDef& graph) override {
return absl::UnimplementedError("Session::Extend()");
}
absl::Status Create(const RunOptions& run_options,
GraphDef&& graph) override {
return absl::UnimplementedError("Session::Create()");
}
absl::Status Extend(const RunOptions& run_options,
GraphDef&& graph) override {
return absl::UnimplementedError("Session::Extend()");
}
absl::Status Close(const RunOptions& run_options) override {
return wrapped_->Close(run_options);
}
absl::Status Run(const RunOptions& run_options,
const std::vector<std::pair<string, Tensor>>& inputs,
const std::vector<string>& output_tensor_names,
const std::vector<string>& target_node_names,
std::vector<Tensor>* outputs,
RunMetadata* run_metadata) override {
return wrapped_->Run(run_options, inputs, output_tensor_names,
target_node_names, outputs, run_metadata);
}
absl::Status Run(
const RunOptions& run_options,
const std::vector<std::pair<std::string, Tensor>>& inputs,
const std::vector<std::string>& output_tensor_names,
const std::vector<std::string>& target_tensor_names,
std::vector<Tensor>* outputs, RunMetadata* run_metadata,
const thread::ThreadPoolOptions& threadpool_options) override {
return wrapped_->Run(run_options, inputs, output_tensor_names,
target_tensor_names, outputs, run_metadata,
threadpool_options);
}
absl::Status PRunSetup(const std::vector<string>& input_names,
const std::vector<string>& output_names,
const std::vector<string>& target_nodes,
string* handle) override {
return absl::UnimplementedError("Session::PRunSetup()");
}
absl::Status PRun(const string& handle,
const std::vector<std::pair<string, Tensor>>& inputs,
const std::vector<string>& output_names,
std::vector<Tensor>* outputs) override {
return absl::UnimplementedError("Session::PRun()");
}
absl::Status ListDevices(std::vector<DeviceAttributes>* response) override {
return wrapped_->ListDevices(response);
}
absl::Status Close() override { return wrapped_->Close(); }
absl::Status LocalDeviceManager(const DeviceMgr** device_mgr) override {
return wrapped_->LocalDeviceManager(device_mgr);
}
absl::Status MakeCallable(const CallableOptions& callable_options,
CallableHandle* out_handle) override {
return wrapped_->MakeCallable(callable_options, out_handle);
}
absl::Status RunCallable(CallableHandle handle,
const std::vector<Tensor>& feed_tensors,
std::vector<Tensor>* fetch_tensors,
RunMetadata* run_metadata) override {
return wrapped_->RunCallable(handle, feed_tensors, fetch_tensors,
run_metadata);
}
absl::Status RunCallable(
CallableHandle handle, const std::vector<Tensor>& feed_tensors,
std::vector<Tensor>* fetch_tensors, RunMetadata* run_metadata,
const thread::ThreadPoolOptions& threadpool_options) override {
return wrapped_->RunCallable(handle, feed_tensors, fetch_tensors,
run_metadata, threadpool_options);
}
absl::Status ReleaseCallable(CallableHandle handle) override {
return wrapped_->ReleaseCallable(handle);
}
absl::Status Finalize() override { return wrapped_->Finalize(); }
private:
const std::unique_ptr<Session> wrapped_;
};
} // namespace
absl::Status LoadSavedModelInternal(const SessionOptions& session_options,
const RunOptions& run_options,
const string& export_dir,
const std::unordered_set<string>& tags,
SavedModelBundleLite* const bundle) {
MetaGraphDef meta_graph_def;
TF_RETURN_IF_ERROR(
ReadMetaGraphDefFromSavedModel(export_dir, tags, &meta_graph_def));
std::unique_ptr<Session> session;
TF_RETURN_IF_ERROR(LoadGraphDefIntoSession(
session_options, std::move(*meta_graph_def.mutable_graph_def()),
&session));
TF_RETURN_IF_ERROR(
RestoreSession(run_options, meta_graph_def, export_dir, &session));
*bundle = SavedModelBundleLite(
std::make_unique<LiteSessionWrapper>(std::move(session)),
std::move(*meta_graph_def.mutable_signature_def()));
return absl::OkStatus();
}
template <typename BundleType>
absl::Status LoadSavedModelGeneric(const SessionOptions& session_options,
const RunOptions& run_options,
const string& export_dir,
const std::unordered_set<string>& tags,
BundleType* const bundle) {
metrics::SavedModelReadApi(kCCLoadLabel).IncrementBy(1);
auto fingerprint_proto =
saved_model::fingerprinting::ReadSavedModelFingerprint(export_dir);
if (fingerprint_proto.ok()) {
// Set gauge cell with saved_model_checksum.
metrics::SavedModelReadFingerprint().Set(
std::to_string(fingerprint_proto->saved_model_checksum()));
}
// TODO(robson): Add tests for the counters.
const uint64 start_microseconds = Env::Default()->NowMicros();
const absl::Status status = LoadSavedModelInternal(
session_options, run_options, export_dir, tags, bundle);
auto log_and_count = [&](const string& status_str) {
LOG(INFO) << "SavedModel load for tags { " << absl::StrJoin(tags, " ")
<< " }; Status: " << status_str << ": " << status << ". Took "
<< GetLatencyMicroseconds(start_microseconds) << " microseconds.";
load_attempt_count->GetCell(export_dir, status_str)->IncrementBy(1);
};
if (status.ok()) {
log_and_count(kLoadAttemptSuccess);
metrics::SavedModelReadPath().Set(export_dir);
} else {
log_and_count(kLoadAttemptFail);
}
load_latency->GetCell(export_dir)
->IncrementBy(GetLatencyMicroseconds(start_microseconds));
return status;
}
absl::Status LoadSavedModel(const SessionOptions& session_options,
const RunOptions& run_options,
const string& export_dir,
const std::unordered_set<string>& tags,
SavedModelBundle* const bundle) {
return LoadSavedModelGeneric<SavedModelBundle>(session_options, run_options,
export_dir, tags, bundle);
}
absl::Status RestoreSession(const RunOptions& run_options,
const MetaGraphDef& meta_graph,
const string& export_dir,
std::unique_ptr<Session>* session) {
const uint64 read_start_microseconds = Env::Default()->NowMicros();
std::vector<AssetFileDef> asset_file_defs;
TF_RETURN_IF_ERROR(internal::GetAssetFileDefs(meta_graph, &asset_file_defs));
if (meta_graph.has_saver_def()) {
TF_RETURN_IF_ERROR(RunRestore(run_options, export_dir,
meta_graph.saver_def().restore_op_name(),
meta_graph.saver_def().filename_tensor_name(),
asset_file_defs, session->get()));
}
// Record walltime spent in restoring graph from disk, but postpone metric
// increments until graph init finishes.
const uint64 restore_graph_walltime =
GetLatencyMicroseconds(read_start_microseconds);
const uint64 graph_init_start_microseconds = Env::Default()->NowMicros();
string init_op_name;
TF_RETURN_IF_ERROR(
internal::GetInitOp(export_dir, meta_graph, &init_op_name));
TF_RETURN_IF_ERROR(RunInitOp(run_options, export_dir, meta_graph,
asset_file_defs, session->get(), init_op_name));
load_latency_by_stage->GetCell(export_dir, "restore_graph")
->Add(restore_graph_walltime);
// Record wall time spent in init op.
load_latency_by_stage->GetCell(export_dir, "init_graph")
->Add(GetLatencyMicroseconds(graph_init_start_microseconds));
return absl::OkStatus();
}
absl::Status LoadSavedModel(const SessionOptions& session_options,
const RunOptions& run_options,
const string& export_dir,
const std::unordered_set<string>& tags,
SavedModelBundleLite* const bundle) {
SessionOptions rewritten_options(session_options);
// We disallow calls to Session::Extend() on the returned session, so we can
// reduce memory consumption by not storing the original GraphDef.
rewritten_options.config.mutable_experimental()
->set_optimize_for_static_graph(true);
// Disallowing the `RunOptions.output_partition_graphs` option (typically used
// in debugging and tests) allows us to reduce memory consumption further by
// not storing the rewritten subgraph for each signature.
rewritten_options.config.mutable_experimental()
->set_disable_output_partition_graphs(true);
// TODO(mrry): Consider specializing the session creation to reduce peak
// RAM consumption by using `Session::Create(GraphDef&&)`.
TF_RETURN_IF_ERROR(LoadSavedModelGeneric(rewritten_options, run_options,
export_dir, tags, bundle));
return absl::OkStatus();
}
bool MaybeSavedModelDirectory(const string& export_dir) {
const string saved_model_pb_path =
io::JoinPath(export_dir, kSavedModelFilenamePb);
const string saved_model_cpb_path =
io::JoinPath(export_dir, kSavedModelFilenameCpb);
const string saved_model_pbtxt_path =
io::JoinPath(export_dir, kSavedModelFilenamePbTxt);
return Env::Default()->FileExists(saved_model_pb_path).ok() ||
Env::Default()->FileExists(saved_model_cpb_path).ok() ||
Env::Default()->FileExists(saved_model_pbtxt_path).ok();
}
} // namespace tensorflow
tensorflow/cc/saved_model/loader.h
+153
View File
@@ -0,0 +1,153 @@
/* Copyright 2016 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.
==============================================================================*/
/// SavedModel loading functions and SavedModelBundle struct.
#ifndef TENSORFLOW_CC_SAVED_MODEL_LOADER_H_
#define TENSORFLOW_CC_SAVED_MODEL_LOADER_H_
#include <string>
#include <unordered_set>
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/public/session.h"
namespace tensorflow {
/// Represents a SavedModel that is loaded from storage.
class SavedModelBundleInterface {
public:
virtual ~SavedModelBundleInterface();
/// Returns the TensorFlow Session that can be used to interact with the
/// SavedModel.
virtual Session* GetSession() const = 0;
/// Returns a map from signature name to SignatureDef for all signatures in
/// in the SavedModel.
virtual const protobuf::Map<std::string, SignatureDef>& GetSignatures()
const = 0;
};
/// SavedModel representation once the SavedModel is loaded from storage.
///
/// NOTE: Prefer to use SavedModelBundleLite in new code, as it consumes less
/// RAM.
struct SavedModelBundle : public SavedModelBundleInterface {
/// A TensorFlow Session does not Close itself on destruction. To avoid
/// resource leaks, we explicitly call Close on Sessions that we create.
~SavedModelBundle() override {
if (session) {
session->Close().IgnoreError();
}
}
SavedModelBundle() = default;
Session* GetSession() const override { return session.get(); }
const protobuf::Map<std::string, SignatureDef>& GetSignatures()
const override {
return meta_graph_def.signature_def();
}
std::unique_ptr<Session> session;
MetaGraphDef meta_graph_def;
std::unique_ptr<GraphDebugInfo> debug_info;
};
// A version of SavedModelBundle that avoids storing a potentially large
// MetaGraphDef. Prefer to use SavedModelBundleLite in new code.
class SavedModelBundleLite : public SavedModelBundleInterface {
public:
SavedModelBundleLite() = default;
SavedModelBundleLite(SavedModelBundleLite&& other) = default;
SavedModelBundleLite& operator=(SavedModelBundleLite&& other) = default;
SavedModelBundleLite(std::unique_ptr<Session> session,
protobuf::Map<std::string, SignatureDef> signatures)
: session_(std::move(session)), signatures_(std::move(signatures)) {}
/// A TensorFlow Session does not Close itself on destruction. To avoid
/// resource leaks, we explicitly call Close on Sessions that we create.
~SavedModelBundleLite() override {
if (session_) {
session_->Close().IgnoreError();
}
}
Session* GetSession() const override { return session_.get(); }
const protobuf::Map<std::string, SignatureDef>& GetSignatures()
const override {
return signatures_;
}
private:
std::unique_ptr<Session> session_;
protobuf::Map<std::string, SignatureDef> signatures_;
};
// Restore variable and resources in the SavedModel export dir for the
// indicated metagraph.
// The recommended way to load a saved model is to call LoadSavedModel,
// which provides an already initialized Metagraph, Session, and DebugInfo.
absl::Status RestoreSession(const RunOptions& run_options,
const MetaGraphDef& meta_graph,
const std::string& export_dir,
std::unique_ptr<Session>* session);
// Initialize a session which wraps this metagraph.
// The recommended way to load a saved model is to call LoadSavedModel,
// which provides an already initialized Metagraph, Session, and DebugInfo.
absl::Status LoadMetagraphIntoSession(const SessionOptions& session_options,
const MetaGraphDef& meta_graph,
std::unique_ptr<Session>* session);
/// Loads a SavedModel from the specified export directory. The MetaGraphDef
/// to be loaded is identified by the supplied tags, corresponding exactly to
/// the set of tags used at SavedModel build time. Stores a SavedModel bundle in
/// *bundle with a session and the requested MetaGraphDef, if found.
///
/// NOTE: Prefer the overload that takes a SavedModelBundleLite* in new code.
absl::Status LoadSavedModel(const SessionOptions& session_options,
const RunOptions& run_options,
const std::string& export_dir,
const std::unordered_set<std::string>& tags,
SavedModelBundle* bundle);
/// Loads a SavedModel from the specified export directory. The MetaGraphDef
/// to be loaded is identified by the supplied tags, corresponding exactly to
/// the set of tags used at SavedModel build time. Stores a SavedModel bundle
/// in *bundle with a session created from the requested MetaGraphDef if found.
///
/// This overload creates a SavedModelBundleLite, which consumes less RAM than
/// an equivalent SavedModelBundle.
absl::Status LoadSavedModel(const SessionOptions& session_options,
const RunOptions& run_options,
const std::string& export_dir,
const std::unordered_set<std::string>& tags,
SavedModelBundleLite* bundle);
/// Checks whether the provided directory could contain a SavedModel. Note that
/// the method does not load any data by itself. If the method returns `false`,
/// the export directory definitely does not contain a SavedModel. If the method
/// returns `true`, the export directory may contain a SavedModel but provides
/// no guarantee that it can be loaded.
bool MaybeSavedModelDirectory(const std::string& export_dir);
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_LOADER_H_
tensorflow/cc/saved_model/loader_util.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2016 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/loader_util.h"
#include <vector>
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/protobuf_internal.h"
namespace tensorflow {
namespace internal {
// A SavedModel may store the name of the initialization op to run in the
// in the SignatureDef (v2) or a collection (v1). If an init_op collection
// exists, then the collection must contain exactly one op.
absl::Status GetInitOp(const string& export_dir,
const MetaGraphDef& meta_graph_def,
string* init_op_name) {
const auto& sig_def_map = meta_graph_def.signature_def();
const auto& init_op_sig_it =
meta_graph_def.signature_def().find(kSavedModelInitOpSignatureKey);
if (init_op_sig_it != sig_def_map.end()) {
const auto& sig_def_outputs = init_op_sig_it->second.outputs();
const auto& sig_def_outputs_it =
sig_def_outputs.find(kSavedModelInitOpSignatureKey);
if (sig_def_outputs_it == sig_def_outputs.end()) {
return absl::FailedPreconditionError(absl::StrCat(
"Could not find output ", kSavedModelInitOpSignatureKey));
}
*init_op_name = sig_def_outputs_it->second.name();
return absl::OkStatus();
}
const auto& collection_def_map = meta_graph_def.collection_def();
string init_op_collection_key;
if (collection_def_map.find(kSavedModelMainOpKey) !=
collection_def_map.end()) {
init_op_collection_key = kSavedModelMainOpKey;
} else {
init_op_collection_key = kSavedModelLegacyInitOpKey;
}
const auto init_op_it = collection_def_map.find(init_op_collection_key);
if (init_op_it != collection_def_map.end()) {
if (init_op_it->second.node_list().value_size() != 1) {
return absl::FailedPreconditionError(
strings::StrCat("Expected exactly one main op in : ", export_dir));
}
*init_op_name = init_op_it->second.node_list().value(0);
}
return absl::OkStatus();
}
absl::Status GetAssetFileDefs(const MetaGraphDef& meta_graph_def,
std::vector<AssetFileDef>* asset_file_defs) {
// With SavedModel v2, we write asset file def into metagraph instead of
// collection, so read from metagraph first.
if (meta_graph_def.asset_file_def_size() > 0) {
for (const auto& asset : meta_graph_def.asset_file_def()) {
asset_file_defs->push_back(asset);
}
return absl::OkStatus();
}
// Fall back to read from collection to be backward compatible with v1.
const auto& collection_def_map = meta_graph_def.collection_def();
const auto assets_it = collection_def_map.find(kSavedModelAssetsKey);
if (assets_it == collection_def_map.end()) {
return absl::OkStatus();
}
const auto& any_assets = assets_it->second.any_list().value();
for (const auto& any_asset : any_assets) {
AssetFileDef asset_file_def;
TF_RETURN_IF_ERROR(
ParseAny(any_asset, &asset_file_def, "tensorflow.AssetFileDef"));
asset_file_defs->push_back(asset_file_def);
}
return absl::OkStatus();
}
} // namespace internal
} // namespace tensorflow
tensorflow/cc/saved_model/loader_util.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_LOADER_UTIL_H_
#define TENSORFLOW_CC_SAVED_MODEL_LOADER_UTIL_H_
#include <string>
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace internal {
// A SavedModel may store the name of the initialization op to run in the
// in the SignatureDef (v2) or a collection (v1). If an init_op collection
// exists, then the collection must contain exactly one op.
absl::Status GetInitOp(const std::string& export_dir,
const MetaGraphDef& meta_graph_def,
std::string* init_op_name);
absl::Status GetAssetFileDefs(const MetaGraphDef& meta_graph_def,
std::vector<AssetFileDef>* asset_file_defs);
} // namespace internal
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_LOADER_UTIL_H_
tensorflow/cc/saved_model/metrics.cc
+324
View File
@@ -0,0 +1,324 @@
/* Copyright 2021 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/metrics.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "json/config.h"
#include "json/json.h"
#include "json/writer.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
#include "tensorflow/core/lib/monitoring/sampler.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
namespace tensorflow {
namespace metrics {
namespace {
// Counter that tracks total number and `write_version` of SavedModels written.
auto* saved_model_write_counter = monitoring::Counter<1>::New(
"/tensorflow/core/saved_model/write/count",
"The number of SavedModels successfully written.", "write_version");
// Counter that tracks total number and `write_version` of SavedModels read.
auto* saved_model_read_counter = monitoring::Counter<1>::New(
"/tensorflow/core/saved_model/read/count",
"The number of SavedModels successfully loaded.", "write_version");
// Counter that tracks number of calls for each SavedModel write API. Summing
// across "api_label" is not expected to equal the ".../write/count" cell value
// because programs can invoke more than one API to save a single SM and
// because the API may error out before successfully writing a SM.
auto* saved_model_write_api = monitoring::Counter<1>::New(
"/tensorflow/core/saved_model/write/api",
"The API used to write the SavedModel.", "api_label");
// Counter that tracks number of calls for each SavedModel read API. Summing
// across "api_label" is not expected to equal the ".../read/count" cell value
// because programs can invoke more than one API to load a single SM and
// because the API may error out before successfully reading a SM.
auto* saved_model_read_api = monitoring::Counter<1>::New(
"/tensorflow/core/saved_model/read/api",
"The API used to load the SavedModel.", "api_label");
// Gauge that contains the fingerprint (saved_model_checksum) of the newly
// written SavedModel.
auto* saved_model_write_fingerprint = monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/write/fingerprint",
"The fingerprint (saved_model_checksum) of the exported SavedModel.");
// Gauge that contains the path (saved_model_path) of the newly written
// SavedModel.
auto* saved_model_write_path = monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/write/path",
"The path (saved_model_path) of the exported SavedModel.");
// Gauge that contains the path (saved_model_path) and the singleprint
// (concatenation of graph_def_program_hash, signature_def_hash,
// saved_object_graph_hash, and checkpoint_hash) of the newly written
// SavedModel.
auto* saved_model_write_path_and_singleprint =
monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/write/path_and_singleprint",
"The path (saved_model_path) and singleprint (concatenation of "
"graph_def_program_hash, signature_def_hash, saved_object_graph_hash, "
"and checkpoint_hash) of the newly written SavedModel.");
// Gauge that contains the fingerprint (saved_model_checksum) of the loaded
// SavedModel.
auto* saved_model_read_fingerprint = monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/read/fingerprint",
"The fingerprint (saved_model_checksum) of the loaded SavedModel.");
// Gauge that contains the path (saved_model_path) of the loaded SavedModel.
auto* saved_model_read_path = monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/read/path",
"The path (saved_model_path) of the loaded SavedModel.");
// Gauge that contains the path (saved_model_path) and the singleprint
// (concatenation of graph_def_program_hash, signature_def_hash,
// saved_object_graph_hash, and checkpoint_hash) of the loaded SavedModel.
auto* saved_model_read_path_and_singleprint =
monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/read/path_and_singleprint",
"The path (saved_model_path) and singleprint (concatenation of "
"graph_def_program_hash, signature_def_hash, saved_object_graph_hash, "
"and checkpoint_hash) of the loaded SavedModel.");
// Gauge that marks whether or not the fingerprint.pb file was found when
// loading the SavedModel.
// Can hold one of the following string values:
// - "FOUND"
// - "NOT_FOUND"
// - "ERROR"
auto* saved_model_found_fingerprint_on_load =
monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/saved_model/found_fingerprint_on_load",
"Whether or not the fingerprint.pb file was found when loading the "
"SavedModel.");
// Distribution of checkpoint write durations.
auto* checkpoint_write_durations = monitoring::Sampler<1>::New(
{
"/tensorflow/core/checkpoint/write/write_durations", // Metric name.
"Distribution of the wall time duration in microseconds of the "
"checkpoint write operation.", // Metric description.
"api_label" // Cell label.
},
// Scale of 1000, growth factor of 1.5 with upper bound of ~184 minutes.
monitoring::Buckets::Exponential(1000, 1.5, 41));
// Distribution of checkpoint read durations.
auto* checkpoint_read_durations = monitoring::Sampler<1>::New(
{
"/tensorflow/core/checkpoint/read/read_durations", // Metric name.
"Distribution of the wall time duration in microseconds of the "
"checkpoint read operation.", // Metric description.
"api_label" // Cell label.
},
// Scale of 1000, growth factor of 1.5 with upper bound of ~184 minutes.
monitoring::Buckets::Exponential(1000, 1.5, 41));
// Distribution of async checkpoint write durations.
auto* async_checkpoint_write_durations = monitoring::Sampler<1>::New(
{
"/tensorflow/core/checkpoint/write/async_write_durations", // Metric
// name.
"Distribution of the wall time duration in microseconds of the async "
"checkpoint write operation", // Metric description.
"api_label" // Cell label.
},
// Scale of 1000, growth factor of 1.5 with upper bound of ~184 minutes.
monitoring::Buckets::Exponential(1000, 1.5, 41));
// Counter that accumulates total time elapsed between module import time and
// the last successful Checkpoint write prior to job preemption or completion.
auto* checkpoint_training_time_saved = monitoring::Counter<1>::New(
"/tensorflow/core/checkpoint/write/training_time_saved",
"Total time in microseconds elapsed between two consecutive write "
"operations in a single job or between Checkpoint construction and the "
"first write operation.",
"api_label");
// Counter that records filesize (MB) of written checkpoint. Contains two cells:
// (api_label, filesize). Cardinality should not be an issue as the filesize
// should be equal among all checkpoints written per job.
auto* checkpoint_size = monitoring::Counter<2>::New(
"/tensorflow/core/checkpoint/write/checkpoint_size",
"Size of checkpoint (.index and sharded data files), rounded to the "
"nearest 100 MB.",
"api_label", "filesize");
} // namespace
// Counter that records how long it took to execute the checkpoint sharding
// callback in microseconds.
auto* sharding_callback_duration = monitoring::Counter<0>::New(
"/tensorflow/core/checkpoint/sharding/callback_duration",
"Sharding callback execution duration in microseconds.");
// Counter that records how many checkpoint shard files were written during
// saving.
auto* num_checkpoint_shards_written = monitoring::Counter<0>::New(
"/tensorflow/core/checkpoint/sharding/num_checkpoint_shards_written",
"Number of checkpoint shard files written during saving.");
// String gauge which describes the callback used to shard the checkpoint during
// saving.
auto* sharding_callback_description = monitoring::Gauge<std::string, 0>::New(
"/tensorflow/core/checkpoint/sharding/callback_description",
"Describes the callback used to shard the checkpoint during saving.");
monitoring::CounterCell& SavedModelWriteCount(absl::string_view write_version) {
return *saved_model_write_counter->GetCell(std::string(write_version));
}
monitoring::CounterCell& SavedModelReadCount(absl::string_view write_version) {
return *saved_model_read_counter->GetCell(std::string(write_version));
}
monitoring::CounterCell& SavedModelWriteApi(absl::string_view api_label) {
return *saved_model_write_api->GetCell(std::string(api_label));
}
monitoring::CounterCell& SavedModelReadApi(absl::string_view api_label) {
return *saved_model_read_api->GetCell(std::string(api_label));
}
monitoring::GaugeCell<std::string>& SavedModelReadFingerprint() {
return *saved_model_read_fingerprint->GetCell();
}
monitoring::GaugeCell<std::string>& SavedModelReadPath() {
return *saved_model_read_path->GetCell();
}
monitoring::GaugeCell<std::string>& SavedModelReadPathAndSingleprint() {
return *saved_model_read_path_and_singleprint->GetCell();
}
monitoring::GaugeCell<std::string>& SavedModelWriteFingerprint() {
return *saved_model_write_fingerprint->GetCell();
}
monitoring::GaugeCell<std::string>& SavedModelWritePath() {
return *saved_model_write_path->GetCell();
}
monitoring::GaugeCell<std::string>& SavedModelWritePathAndSingleprint() {
return *saved_model_write_path_and_singleprint->GetCell();
}
std::string MakeFingerprintJson(FingerprintDef fingerprint_def) {
Json::Value fingerprint = Json::objectValue;
fingerprint["saved_model_checksum"] =
Json::UInt64(fingerprint_def.saved_model_checksum());
fingerprint["graph_def_program_hash"] =
Json::UInt64(fingerprint_def.graph_def_program_hash());
fingerprint["signature_def_hash"] =
Json::UInt64(fingerprint_def.signature_def_hash());
fingerprint["saved_object_graph_hash"] =
Json::UInt64(fingerprint_def.saved_object_graph_hash());
fingerprint["checkpoint_hash"] =
Json::UInt64(fingerprint_def.checkpoint_hash());
Json::StreamWriterBuilder json_factory;
return Json::writeString(json_factory, fingerprint);
}
absl::StatusOr<std::string> MakeSavedModelPathAndSingleprint(
std::string path, std::string singleprint) {
if (path.empty()) {
return absl::InvalidArgumentError(
"Invalid path_and_singleprint argument. Empty path.");
}
if (singleprint.empty()) {
return absl::InvalidArgumentError(
"Invalid path_and_singleprint argument. Empty singleprint.");
}
return absl::StrCat(path, ":", singleprint);
}
absl::StatusOr<std::pair<std::string, std::string>>
ParseSavedModelPathAndSingleprint(std::string path_and_singleprint) {
size_t delimiter = path_and_singleprint.rfind(':');
if (delimiter == std::string::npos) {
return absl::InvalidArgumentError(
"Invalid path_and_singleprint argument. Found no delimeter.");
}
std::string path = path_and_singleprint.substr(0, delimiter);
if (path.empty()) {
return absl::InvalidArgumentError(
"Invalid path_and_singleprint argument. Empty path.");
}
std::string singleprint = path_and_singleprint.substr(delimiter + 1);
if (singleprint.empty()) {
return absl::InvalidArgumentError(
"Invalid path_and_singleprint argument. Empty singleprint.");
}
return std::pair<std::string, std::string>(path, singleprint);
}
monitoring::GaugeCell<std::string>& SavedModelFoundFingerprintOnLoad() {
return *saved_model_found_fingerprint_on_load->GetCell();
}
monitoring::SamplerCell& CheckpointReadDuration(absl::string_view api_label) {
return *checkpoint_read_durations->GetCell(std::string(api_label));
}
monitoring::SamplerCell& CheckpointWriteDuration(absl::string_view api_label) {
return *checkpoint_write_durations->GetCell(std::string(api_label));
}
monitoring::SamplerCell& AsyncCheckpointWriteDuration(
absl::string_view api_label) {
return *async_checkpoint_write_durations->GetCell(std::string(api_label));
}
monitoring::CounterCell& TrainingTimeSaved(absl::string_view api_label) {
return *checkpoint_training_time_saved->GetCell(std::string(api_label));
}
monitoring::CounterCell& CheckpointSize(absl::string_view api_label,
int64_t filesize) {
return *checkpoint_size->GetCell(std::string(api_label),
std::to_string(filesize));
}
monitoring::CounterCell& ShardingCallbackDuration() {
return *sharding_callback_duration->GetCell();
}
monitoring::CounterCell& NumCheckpointShardsWritten() {
return *num_checkpoint_shards_written->GetCell();
}
monitoring::GaugeCell<std::string>& ShardingCallbackDescription() {
return *sharding_callback_description->GetCell();
}
} // namespace metrics
} // namespace tensorflow
tensorflow/cc/saved_model/metrics.h
+147
View File
@@ -0,0 +1,147 @@
/* Copyright 2021 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.
==============================================================================*/
// APIs for accessing SavedModel and checkpoint metric objects.
//
// In order to collect the data from these metrics, please add the metrics to
// the provided monitoring platform. Unless configured with a user-specified
// monitoring platform, the data is not collected in OSS.
#ifndef TENSORFLOW_CC_SAVED_MODEL_METRICS_H_
#define TENSORFLOW_CC_SAVED_MODEL_METRICS_H_
#include <cstdint>
#include <string>
#include <utility>
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/lib/monitoring/counter.h"
#include "tensorflow/core/lib/monitoring/gauge.h"
#include "tensorflow/core/lib/monitoring/sampler.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
namespace tensorflow {
namespace metrics {
const char kFingerprintFound[] = "FOUND";
const char kFingerprintNotFound[] = "NOT_FOUND";
const char kFingerprintError[] = "ERROR";
// Returns "/tensorflow/core/saved_model/write/count" cell. This metric
// has 1 field "write_version", which is equal to the
// `tensorflow::libexport::GetWriteVersion` of the protobuf and should be
// incremented when a SavedModel has been successfully written.
monitoring::CounterCell& SavedModelWriteCount(absl::string_view write_version);
// Returns "/tensorflow/core/saved_model/read/count" cell. This metric
// has 1 field "write_version", which is equal to the
// `tensorflow::libexport::GetWriteVersion` of the protobuf, and should be
// incremented when a SavedModel has been successfully read.
monitoring::CounterCell& SavedModelReadCount(absl::string_view write_version);
// Returns "/tensorflow/core/saved_model/write/api" cell. This metric has 1
// field "api_label" which corresponds to a SavedModel write API. The cell for
// `foo` should be incremented when the write API `foo` is called.
monitoring::CounterCell& SavedModelWriteApi(absl::string_view api_label);
// Returns "/tensorflow/core/saved_model/read/api" cell. This metric has 1
// field "api_label" which corresponds to a SavedModel read API. The cell for
// `foo` should be incremented when the read API `foo` is called.
monitoring::CounterCell& SavedModelReadApi(absl::string_view api_label);
// Returns "/tensorflow/core/saved_model/write/fingerprint" cell, which contains
// the saved_model_checksum of the SM's fingerprint when it is exported.
monitoring::GaugeCell<std::string>& SavedModelWriteFingerprint();
// Returns "/tensorflow/core/saved_model/write/path" cell, which contains
// the saved_model_path of the SM when it is exported.
monitoring::GaugeCell<std::string>& SavedModelWritePath();
// Returns "/tensorflow/core/saved_model/write/path_and_fingerprint" cell, which
// contains the path (saved_model_path) and fingerprint (concatenation of
// graph_def_program_hash, signature_def_hash, saved_object_graph_hash,
// and checkpoint_hash) of the SavedModel when it is exported.
monitoring::GaugeCell<std::string>& SavedModelWritePathAndSingleprint();
// Returns "/tensorflow/core/saved_model/read/fingerprint" cell, wich contains
// the saved_model_checksum of the SM's fingerprint when it is imported.
monitoring::GaugeCell<std::string>& SavedModelReadFingerprint();
// Returns "/tensorflow/core/saved_model/read/path" cell, wich contains
// the saved_model_path of the SM when it is imported.
monitoring::GaugeCell<std::string>& SavedModelReadPath();
// Returns "/tensorflow/core/saved_model/read/path_and_fingerprint" cell, which
// contains the path (saved_model_path) and singleprint (concatenation of
// graph_def_program_hash, signature_def_hash, saved_object_graph_hash,
// and checkpoint_hash) of the SavedModel when it is imported.
monitoring::GaugeCell<std::string>& SavedModelReadPathAndSingleprint();
// Returns the fingerprint as a Json string.
std::string MakeFingerprintJson(FingerprintDef fingerprint_def);
// Returns canonical string concatenation of path and singleprint.
absl::StatusOr<std::string> MakeSavedModelPathAndSingleprint(
std::string path, std::string singleprint);
// Returns path and singleprint as a pair, parsed canonically from the string
// metric.
absl::StatusOr<std::pair<std::string, std::string>>
ParseSavedModelPathAndSingleprint(std::string path_and_singleprint);
// Returns string status indicating whether or not the fingerprint.pb file was
// found when loading the SavedModel.
monitoring::GaugeCell<std::string>& SavedModelFoundFingerprintOnLoad();
// Returns "/tensorflow/core/checkpoint/read/read_durations" cell belonging to
// field `api_label`.
monitoring::SamplerCell& CheckpointReadDuration(absl::string_view api_label);
// Returns "/tensorflow/core/checkpoint/write/write_durations" cell belonging to
// field `api_label`.
monitoring::SamplerCell& CheckpointWriteDuration(absl::string_view api_label);
// Returns "/tensorflow/core/checkpoint/write/async_write_durations" cell
// belonging to field `api_label`.
monitoring::SamplerCell& AsyncCheckpointWriteDuration(
absl::string_view api_label);
// Returns "/tensorflow/core/checkpoint/write/training_time_saved" cell
// belonging to field `api_label`.
monitoring::CounterCell& TrainingTimeSaved(absl::string_view api_label);
// Returns "/tensorflow/core/checkpoint/write/checkpoint_size" cell
// belonging to field (`api_label`, `filesize`).
monitoring::CounterCell& CheckpointSize(absl::string_view api_label,
int64_t filesize);
// Returns "/tensorflow/core/checkpoint/sharding/callback_duration" cell which
// describes how long it took to execute the checkpoint sharding callback in
// microseconds.
monitoring::CounterCell& ShardingCallbackDuration();
// Returns "/tensorflow/core/checkpoint/sharding/num_checkpoint_shards_written"
// cell which describes how many checkpoint shard files were written during
// saving.
monitoring::CounterCell& NumCheckpointShardsWritten();
// Returns "/tensorflow/core/checkpoint/sharding/callback_description" cell
// which describes the callback used to shard the checkpoint during saving.
monitoring::GaugeCell<std::string>& ShardingCallbackDescription();
} // namespace metrics
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_METRICS_H_
tensorflow/cc/saved_model/metrics_test.cc
+208
View File
@@ -0,0 +1,208 @@
/* Copyright 2021 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/metrics.h"
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "json/json.h"
#include "json/reader.h"
#include "tensorflow/core/platform/status_matchers.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/fingerprint.pb.h"
#include "tsl/platform/statusor.h"
namespace tensorflow {
namespace metrics {
// The value of the cells for each metric persists across tests.
TEST(MetricsTest, TestSavedModelWrite) {
EXPECT_EQ(SavedModelWriteApi("foo").value(), 0);
SavedModelWriteApi("foo").IncrementBy(1);
EXPECT_EQ(SavedModelWriteApi("foo").value(), 1);
EXPECT_EQ(SavedModelWriteCount("1").value(), 0);
SavedModelWriteCount("1").IncrementBy(1);
EXPECT_EQ(SavedModelWriteCount("1").value(), 1);
}
TEST(MetricsTest, TestSavedModelRead) {
SavedModelReadApi("bar").IncrementBy(1);
EXPECT_EQ(SavedModelReadApi("bar").value(), 1);
SavedModelReadCount("2").IncrementBy(1);
EXPECT_EQ(SavedModelReadCount("2").value(), 1);
SavedModelReadApi("baz").IncrementBy(1);
EXPECT_EQ(SavedModelReadApi("baz").value(), 1);
SavedModelReadCount("2").IncrementBy(1);
EXPECT_EQ(SavedModelReadCount("2").value(), 2);
}
TEST(MetricsTest, TestCheckpointRead) {
EXPECT_EQ(CheckpointReadDuration("foo").value().num(), 0);
CheckpointReadDuration("foo").Add(100);
EXPECT_EQ(CheckpointReadDuration("foo").value().num(), 1);
}
TEST(MetricsTest, TestCheckpointWrite) {
EXPECT_EQ(CheckpointWriteDuration("foo").value().num(), 0);
CheckpointWriteDuration("foo").Add(100);
EXPECT_EQ(CheckpointWriteDuration("foo").value().num(), 1);
}
TEST(MetricsTest, TestAsyncCheckpointWrite) {
EXPECT_EQ(AsyncCheckpointWriteDuration("foo").value().num(), 0);
AsyncCheckpointWriteDuration("foo").Add(100);
EXPECT_EQ(AsyncCheckpointWriteDuration("foo").value().num(), 1);
}
TEST(MetricsTest, TestTrainingTimeSaved) {
EXPECT_EQ(TrainingTimeSaved("foo").value(), 0);
TrainingTimeSaved("foo").IncrementBy(100);
EXPECT_EQ(TrainingTimeSaved("foo").value(), 100);
}
TEST(MetricsTest, TestCheckpointSize) {
EXPECT_EQ(CheckpointSize("foo", 10).value(), 0);
CheckpointSize("foo", 10).IncrementBy(1);
EXPECT_EQ(CheckpointSize("foo", 10).value(), 1);
}
TEST(MetricsTest, TestWriteFingerprint) {
EXPECT_EQ(SavedModelWriteFingerprint().value(), "");
SavedModelWriteFingerprint().Set("foo");
EXPECT_EQ(SavedModelWriteFingerprint().value(), "foo");
SavedModelWriteFingerprint().Set("bar");
EXPECT_EQ(SavedModelWriteFingerprint().value(), "bar");
}
TEST(MetricsTest, TestWritePath) {
EXPECT_EQ(SavedModelWritePath().value(), "");
SavedModelWritePath().Set("foo");
EXPECT_EQ(SavedModelWritePath().value(), "foo");
SavedModelWritePath().Set("bar");
EXPECT_EQ(SavedModelWritePath().value(), "bar");
}
TEST(MetricsTest, TestWritePathAndSingleprint) {
EXPECT_EQ(SavedModelWritePathAndSingleprint().value(), "");
SavedModelWritePathAndSingleprint().Set("foo");
EXPECT_EQ(SavedModelWritePathAndSingleprint().value(), "foo");
SavedModelWritePathAndSingleprint().Set("bar");
EXPECT_EQ(SavedModelWritePathAndSingleprint().value(), "bar");
EXPECT_EQ(
MakeSavedModelPathAndSingleprint("path", "singleprint").value_or(""),
"path:singleprint");
}
TEST(MetricsTest, TestInvalidMakePathAndSingleprint) {
EXPECT_THAT(MakeSavedModelPathAndSingleprint("", "singleprint"),
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
EXPECT_THAT(MakeSavedModelPathAndSingleprint("path", ""),
absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
}
TEST(MetricsTest, TestReadFingerprint) {
EXPECT_EQ(SavedModelReadFingerprint().value(), "");
SavedModelReadFingerprint().Set("foo");
EXPECT_EQ(SavedModelReadFingerprint().value(), "foo");
SavedModelReadFingerprint().Set("bar");
EXPECT_EQ(SavedModelReadFingerprint().value(), "bar");
}
TEST(MetricsTest, TestReadPath) {
EXPECT_EQ(SavedModelReadPath().value(), "");
SavedModelReadPath().Set("foo");
EXPECT_EQ(SavedModelReadPath().value(), "foo");
SavedModelReadPath().Set("bar");
EXPECT_EQ(SavedModelReadPath().value(), "bar");
}
TEST(MetricsTest, TestReadPathAndSingleprint) {
EXPECT_EQ(SavedModelReadPathAndSingleprint().value(), "");
SavedModelReadPathAndSingleprint().Set("foo");
EXPECT_EQ(SavedModelReadPathAndSingleprint().value(), "foo");
SavedModelReadPathAndSingleprint().Set("bar");
EXPECT_EQ(SavedModelReadPathAndSingleprint().value(), "bar");
TF_ASSERT_OK_AND_ASSIGN(
auto path_singleprint,
ParseSavedModelPathAndSingleprint("path/model:name:singleprint"));
auto [path, singleprint] = path_singleprint;
EXPECT_EQ(path, "path/model:name");
EXPECT_EQ(singleprint, "singleprint");
}
TEST(MetricsTest, TestMakeFingerprintJson) {
FingerprintDef fingerprint;
fingerprint.set_saved_model_checksum(1);
fingerprint.set_graph_def_program_hash(2);
fingerprint.set_signature_def_hash(3);
fingerprint.set_saved_object_graph_hash(4);
fingerprint.set_checkpoint_hash(5);
std::string serialized_fingerprint_json = MakeFingerprintJson(fingerprint);
EXPECT_EQ(
serialized_fingerprint_json,
"{\n\t\"checkpoint_hash\" : 5,\n\t\"graph_def_program_hash\" : "
"2,\n\t\"saved_model_checksum\" : 1,\n\t\"saved_object_graph_hash\" : "
"4,\n\t\"signature_def_hash\" : 3\n}");
Json::Value fingerprint_json = Json::objectValue;
Json::Reader reader = Json::Reader();
reader.parse(serialized_fingerprint_json, fingerprint_json);
EXPECT_EQ(fingerprint_json["saved_model_checksum"].asUInt64(), 1);
EXPECT_EQ(fingerprint_json["graph_def_program_hash"].asUInt64(), 2);
EXPECT_EQ(fingerprint_json["signature_def_hash"].asUInt64(), 3);
EXPECT_EQ(fingerprint_json["saved_object_graph_hash"].asUInt64(), 4);
EXPECT_EQ(fingerprint_json["checkpoint_hash"].asUInt64(), 5);
}
TEST(MetricsTest, TestFoundFingerprintOnLoad) {
EXPECT_EQ(SavedModelFoundFingerprintOnLoad().value(), "");
SavedModelFoundFingerprintOnLoad().Set(kFingerprintFound);
EXPECT_EQ(SavedModelFoundFingerprintOnLoad().value(), "FOUND");
SavedModelFoundFingerprintOnLoad().Set(kFingerprintNotFound);
EXPECT_EQ(SavedModelFoundFingerprintOnLoad().value(), "NOT_FOUND");
SavedModelFoundFingerprintOnLoad().Set(kFingerprintError);
EXPECT_EQ(SavedModelFoundFingerprintOnLoad().value(), "ERROR");
}
TEST(MetricsTest, TestShardingCallbackDuration) {
EXPECT_EQ(ShardingCallbackDuration().value(), 0);
ShardingCallbackDuration().IncrementBy(100);
EXPECT_EQ(ShardingCallbackDuration().value(), 100);
}
TEST(MetricsTest, TestNumCheckpointShardsWritten) {
EXPECT_EQ(NumCheckpointShardsWritten().value(), 0);
NumCheckpointShardsWritten().IncrementBy(10);
EXPECT_EQ(NumCheckpointShardsWritten().value(), 10);
}
TEST(MetricsTest, TestShardingCallbackDescription) {
EXPECT_EQ(ShardingCallbackDescription().value(), "");
ShardingCallbackDescription().Set("foo");
EXPECT_EQ(ShardingCallbackDescription().value(), "foo");
}
} // namespace metrics
} // namespace tensorflow
tensorflow/cc/saved_model/python/BUILD
+18
View File
@@ -0,0 +1,18 @@
# Description:
# CLIF wrappers for TensorFlow SavedModels.
load("//tensorflow/core/platform:build_config.bzl", "tf_py_clif_cc")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
tf_py_clif_cc(
name = "loader",
srcs = ["loader.clif"],
deps = [
"//tensorflow/cc/saved_model:loader",
],
)
tensorflow/cc/saved_model/python/loader.clif
+19
View File
@@ -0,0 +1,19 @@
# Copyright 2026 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.
# ==============================================================================
from "third_party/tensorflow/cc/saved_model/loader.h":
namespace `tensorflow`:
class SavedModelBundle:
def __init__(self)
tensorflow/cc/saved_model/reader.cc
+163
View File
@@ -0,0 +1,163 @@
/* Copyright 2018 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/reader.h"
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include "absl/memory/memory.h"
#include "absl/status/statusor.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/statusor.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/util/tensor_bundle/byte_swap_tensor.h"
// Placeholder for protosplitter merger include.
#define IS_OSS true
namespace tensorflow {
absl::StatusOr<MetaGraphDef*> FindMetaGraphDef(
const std::unordered_set<string>& tags, SavedModel* saved_model_proto) {
LOG(INFO) << "Reading meta graph with tags { " << absl::StrJoin(tags, " ")
<< " }";
for (MetaGraphDef& graph_def : *saved_model_proto->mutable_meta_graphs()) {
// Get tags from the graph_def.
std::unordered_set<string> graph_tags;
for (const string& tag : graph_def.meta_info_def().tags()) {
graph_tags.insert(tag);
}
// Match with the set of tags provided.
if (graph_tags == tags) {
MetaGraphDef* meta_graph_def = &graph_def;
// Correct the endiness of Tensor content on big-endian system
if (!port::kLittleEndian) {
TF_RETURN_IF_ERROR(ByteSwapTensorContentInMetaGraphDef(meta_graph_def));
}
return meta_graph_def;
}
}
return Status(
absl::StatusCode::kNotFound,
absl::StrCat(
"Could not find meta graph def matching supplied tags: { ",
absl::StrJoin(tags, " "),
" }. To inspect available tag-sets in the SavedModel, please "
"use the SavedModel CLI: `saved_model_cli`"));
}
// Reads the SavedModel proto from saved_model.pb in `export_dir`.
// Returns a failure status when the SavedModel file does not exist.
Status ReadSavedModel(absl::string_view export_dir,
SavedModel* saved_model_proto) {
LOG(INFO) << "Reading SavedModel from: " << export_dir;
if (IS_OSS) {
const std::string saved_model_pb_path =
io::JoinPath(export_dir, kSavedModelFilenamePb);
TF_ASSIGN_OR_RETURN(
bool saved_model_pb_exists,
internal::FileExists(Env::Default(), saved_model_pb_path));
if (saved_model_pb_exists) {
Status result = ReadBinaryProto(Env::Default(), saved_model_pb_path,
saved_model_proto);
if (result.ok()) {
metrics::SavedModelReadCount(
saved_model::GetWriteVersion(*saved_model_proto))
.IncrementBy(1);
}
return result;
}
}
const std::string saved_model_pbtxt_path =
io::JoinPath(export_dir, kSavedModelFilenamePbTxt);
auto saved_model_pbtxt_exists =
internal::FileExists(Env::Default(), saved_model_pbtxt_path);
if (saved_model_pbtxt_exists.value_or(false)) {
absl::Status result = ReadTextProto(Env::Default(), saved_model_pbtxt_path,
saved_model_proto);
if (result.ok()) {
metrics::SavedModelReadCount(
saved_model::GetWriteVersion(*saved_model_proto))
.IncrementBy(1);
}
return result;
}
if (!IS_OSS) {
// Only use Merger outside of OSS.
// Placeholder for protosplitter merger call.
}
return absl::Status(
absl::StatusCode::kNotFound,
absl::StrCat("Could not find SavedModel .pb or .pbtxt at supplied "
"export directory path: ",
export_dir,
". Check that "
"the directory exists and that you have the right "
"permissions for accessing it."));
}
Status ReadMetaGraphDefFromSavedModel(absl::string_view export_dir,
const std::unordered_set<string>& tags,
MetaGraphDef* const meta_graph_def) {
SavedModel saved_model_proto;
TF_RETURN_IF_ERROR(ReadSavedModel(export_dir, &saved_model_proto));
TF_ASSIGN_OR_RETURN(MetaGraphDef * m,
FindMetaGraphDef(tags, &saved_model_proto));
*meta_graph_def = std::move(*m);
return absl::OkStatus();
}
absl::Status ReadSavedModelDebugInfoIfPresent(
absl::string_view export_dir,
std::unique_ptr<GraphDebugInfo>* debug_info_proto) {
LOG(INFO) << "Reading SavedModel debug info (if present) from: "
<< export_dir;
const string debug_info_pb_path =
io::JoinPath(export_dir, "debug", "saved_model_debug_info.pb");
TF_ASSIGN_OR_RETURN(bool debug_info_pb_exists,
internal::FileExists(Env::Default(), debug_info_pb_path));
if (debug_info_pb_exists) {
GraphDebugInfo debug_info;
TF_RETURN_IF_ERROR(
ReadBinaryProto(Env::Default(), debug_info_pb_path, &debug_info));
*debug_info_proto = std::make_unique<GraphDebugInfo>(std::move(debug_info));
}
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/cc/saved_model/reader.h
+59
View File
@@ -0,0 +1,59 @@
/* Copyright 2018 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.
==============================================================================*/
/// Functions to read the SavedModel proto, or parts of it.
#ifndef TENSORFLOW_CC_SAVED_MODEL_READER_H_
#define TENSORFLOW_CC_SAVED_MODEL_READER_H_
#include <memory>
#include <unordered_set>
#include "absl/status/statusor.h"
#include "tensorflow/core/framework/graph_debug_info.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
namespace tensorflow {
absl::Status ReadSavedModel(absl::string_view export_dir,
SavedModel* saved_model_proto);
// Finds and returns the MetaGraphDef (within the provided SavedModel) that
// matches the given set of tags. The lifetime of the returned MetaGraphDef is
// the same as the lifetime of `saved_model_proto`.
//
// FindMetaGraphDef returns a failure status when no MetaGraphDef matches the
// provided tags.
absl::StatusOr<MetaGraphDef*> FindMetaGraphDef(
const std::unordered_set<std::string>& tags, SavedModel* saved_model_proto);
// Reads the SavedModel proto from saved_model.pb(txt) in the given directory,
// finds the MetaGraphDef that matches the given set of tags and writes it to
// the `meta_graph_def` parameter. Returns a failure status when the SavedModel
// file does not exist or no MetaGraphDef matches the tags.
absl::Status ReadMetaGraphDefFromSavedModel(
absl::string_view export_dir, const std::unordered_set<std::string>& tags,
MetaGraphDef* meta_graph_def);
// Store debug info from the SavedModel export dir.
absl::Status ReadSavedModelDebugInfoIfPresent(
absl::string_view export_dir,
std::unique_ptr<GraphDebugInfo>* debug_info_proto);
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_READER_H_
tensorflow/cc/saved_model/reader_test.cc
+143
View File
@@ -0,0 +1,143 @@
/* Copyright 2018 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/reader.h"
#include <gmock/gmock.h>
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/tag_constants.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/resource_loader.h"
namespace tensorflow {
namespace {
string TestDataPbTxt() {
return io::JoinPath("tensorflow", "cc", "saved_model", "testdata",
"half_plus_two_pbtxt", "00000123");
}
string TestDataSharded() {
return io::JoinPath("tensorflow", "cc", "saved_model", "testdata",
"half_plus_two", "00000123");
}
string ChunkedSavedModel() {
return io::JoinPath("tensorflow", "cc", "saved_model", "testdata",
"chunked_saved_model", "chunked_model");
}
string NonChunkedSavedModel() {
return io::JoinPath("tensorflow", "cc", "saved_model", "testdata",
"chunked_saved_model", "non_chunked_model");
}
class ReaderTest : public ::testing::Test {
protected:
ReaderTest() {}
void CheckMetaGraphDef(const MetaGraphDef& meta_graph_def) {
const auto& tags = meta_graph_def.meta_info_def().tags();
EXPECT_TRUE(std::find(tags.begin(), tags.end(), kSavedModelTagServe) !=
tags.end());
EXPECT_NE(meta_graph_def.meta_info_def().tensorflow_version(), "");
EXPECT_EQ(
meta_graph_def.signature_def().at("serving_default").method_name(),
"tensorflow/serving/predict");
}
};
TEST_F(ReaderTest, TagMatch) {
MetaGraphDef meta_graph_def;
const string export_dir = GetDataDependencyFilepath(TestDataSharded());
TF_ASSERT_OK(ReadMetaGraphDefFromSavedModel(export_dir, {kSavedModelTagServe},
&meta_graph_def));
CheckMetaGraphDef(meta_graph_def);
}
TEST_F(ReaderTest, NoTagMatch) {
MetaGraphDef meta_graph_def;
const string export_dir = GetDataDependencyFilepath(TestDataSharded());
absl::Status st = ReadMetaGraphDefFromSavedModel(export_dir, {"missing-tag"},
&meta_graph_def);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(),
"Could not find meta graph def matching supplied tags: { missing-tag }"))
<< st.message();
}
TEST_F(ReaderTest, NoTagMatchMultiple) {
MetaGraphDef meta_graph_def;
const string export_dir = GetDataDependencyFilepath(TestDataSharded());
absl::Status st = ReadMetaGraphDefFromSavedModel(
export_dir, {kSavedModelTagServe, "missing-tag"}, &meta_graph_def);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(), "Could not find meta graph def matching supplied tags: "))
<< st.message();
}
TEST_F(ReaderTest, InvalidExportPath) {
MetaGraphDef meta_graph_def;
const string export_dir = GetDataDependencyFilepath("missing-path");
absl::Status st = ReadMetaGraphDefFromSavedModel(
export_dir, {kSavedModelTagServe}, &meta_graph_def);
EXPECT_FALSE(st.ok());
}
TEST_F(ReaderTest, ReadSavedModelDebugInfoIfPresent) {
const string export_dir = GetDataDependencyFilepath(TestDataSharded());
std::unique_ptr<GraphDebugInfo> debug_info_proto;
TF_ASSERT_OK(ReadSavedModelDebugInfoIfPresent(export_dir, &debug_info_proto));
}
TEST_F(ReaderTest, MetricsNotUpdatedFailedRead) {
MetaGraphDef meta_graph_def;
const int read_count_v1 = metrics::SavedModelReadCount("1").value();
const int read_count_v2 = metrics::SavedModelReadCount("2").value();
const string export_dir = GetDataDependencyFilepath("missing-path");
absl::Status st =
ReadMetaGraphDefFromSavedModel(export_dir, {"serve"}, &meta_graph_def);
EXPECT_FALSE(st.ok());
EXPECT_EQ(metrics::SavedModelReadCount("1").value(), read_count_v1);
EXPECT_EQ(metrics::SavedModelReadCount("2").value(), read_count_v2);
}
TEST_F(ReaderTest, MetricsUpdatedSuccessfulRead) {
MetaGraphDef meta_graph_def;
const int read_count_v1 = metrics::SavedModelReadCount("1").value();
const string export_dir = GetDataDependencyFilepath(TestDataSharded());
absl::Status st =
ReadMetaGraphDefFromSavedModel(export_dir, {"serve"}, &meta_graph_def);
EXPECT_EQ(metrics::SavedModelReadCount("1").value(), read_count_v1 + 1);
}
// Placeholder for protosplitter merger merge test.
} // namespace
} // namespace tensorflow
tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc
+266
View File
@@ -0,0 +1,266 @@
/* Copyright 2016 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.
==============================================================================*/
#include "absl/status/status.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/signature_constants.h"
#include "tensorflow/cc/saved_model/tag_constants.h"
#include "tensorflow/core/example/example.pb.h"
#include "tensorflow/core/example/feature.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace {
constexpr char kTestDataPbTxt[] =
"cc/saved_model/testdata/half_plus_two_pbtxt/00000123";
constexpr char kTestDataMainOp[] =
"cc/saved_model/testdata/half_plus_two_main_op/00000123";
constexpr char kTestDataSharded[] =
"cc/saved_model/testdata/half_plus_two/00000123";
constexpr char kTestDataInitOpV2[] =
"cc/saved_model/testdata/half_plus_two_v2/00000123";
class LoaderTest : public ::testing::Test {
protected:
LoaderTest() {}
string MakeSerializedExample(float x) {
tensorflow::Example example;
auto* feature_map = example.mutable_features()->mutable_feature();
(*feature_map)["x"].mutable_float_list()->add_value(x);
return example.SerializeAsString();
}
void ValidateAssets(const string& export_dir,
const SavedModelBundleLite& bundle) {
const string asset_directory =
io::JoinPath(export_dir, kSavedModelAssetsDirectory);
const string asset_filename = "foo.txt";
const string asset_filepath = io::JoinPath(asset_directory, asset_filename);
TF_EXPECT_OK(Env::Default()->FileExists(asset_filepath));
std::vector<Tensor> path_outputs;
TF_ASSERT_OK(
bundle.GetSession()->Run({}, {"filename_tensor:0"}, {}, &path_outputs));
ASSERT_EQ(1, path_outputs.size());
test::ExpectTensorEqual<tstring>(
test::AsTensor<tstring>({"foo.txt"}, TensorShape({})), path_outputs[0]);
}
void CheckSavedModelBundle(const string& export_dir,
const SavedModelBundleLite& bundle) {
ValidateAssets(export_dir, bundle);
// Retrieve the regression signature from the bundle.
const auto& signature_def = bundle.GetSignatures().at("regress_x_to_y");
const string input_name = signature_def.inputs().at(kRegressInputs).name();
const string output_name =
signature_def.outputs().at(kRegressOutputs).name();
std::vector<tstring> serialized_examples;
for (float x : {0, 1, 2, 3}) {
serialized_examples.push_back(MakeSerializedExample(x));
}
// Validate the half plus two behavior.
Tensor input =
test::AsTensor<tstring>(serialized_examples, TensorShape({4}));
std::vector<Tensor> outputs;
TF_ASSERT_OK(bundle.GetSession()->Run({{input_name, input}}, {output_name},
{}, &outputs));
ASSERT_EQ(outputs.size(), 1);
test::ExpectTensorEqual<float>(
outputs[0],
test::AsTensor<float>({2, 2.5, 3, 3.5}, TensorShape({4, 1})));
// Validate the `output_partition_graphs` is not supported.
RunOptions run_options;
run_options.set_output_partition_graphs(true);
RunMetadata run_metadata;
absl::Status s =
bundle.GetSession()->Run(run_options, {{input_name, input}},
{output_name}, {}, &outputs, &run_metadata);
ASSERT_TRUE(absl::IsInvalidArgument(s));
}
};
// Test for resource leaks related to TensorFlow session closing requirements
// when loading and unloading large numbers of SavedModelBundles.
// TODO(sukritiramesh): Increase run iterations and move outside of the test
// suite.
TEST_F(LoaderTest, ResourceLeakTest) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
for (int i = 0; i < 100; ++i) {
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
}
TEST_F(LoaderTest, ExtendFailsTest) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
absl::Status s = bundle.GetSession()->Extend({});
ASSERT_TRUE(absl::IsUnimplemented(s));
}
TEST_F(LoaderTest, TagMatch) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, NoTagMatch) {
SavedModelBundleLite bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{"missing-tag"}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(),
"Could not find meta graph def matching supplied tags: { missing-tag }"))
<< st.message();
}
TEST_F(LoaderTest, NoTagMatchMultiple) {
SavedModelBundleLite bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st =
LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe, "missing-tag"}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(), "Could not find meta graph def matching supplied tags: "))
<< st.message();
}
TEST_F(LoaderTest, SessionCreationFailure) {
SavedModelBundleLite bundle;
// Use invalid SessionOptions to cause session creation to fail. Default
// options work, so provide an invalid value for the target field.
SessionOptions session_options;
constexpr char kInvalidTarget[] = "invalid target";
session_options.target = kInvalidTarget;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(st.message(), kInvalidTarget)) << st.message();
}
TEST_F(LoaderTest, PbtxtFormat) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataPbTxt);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, MainOpFormat) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataMainOp);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, InvalidExportPath) {
SavedModelBundleLite bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "missing-path");
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
}
TEST_F(LoaderTest, MaybeSavedModelDirectory) {
// Valid SavedModel directory.
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
EXPECT_TRUE(MaybeSavedModelDirectory(export_dir));
// Directory that does not exist.
const string missing_export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "missing-path");
EXPECT_FALSE(MaybeSavedModelDirectory(missing_export_dir));
// Directory that exists but is an invalid SavedModel location.
const string invalid_export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model");
EXPECT_FALSE(MaybeSavedModelDirectory(invalid_export_dir));
}
TEST_F(LoaderTest, SavedModelInitOpV2Format) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataInitOpV2);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
} // namespace
} // namespace tensorflow
tensorflow/cc/saved_model/saved_model_bundle_test.cc
+405
View File
@@ -0,0 +1,405 @@
/* Copyright 2016 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.
==============================================================================*/
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/metrics.h"
#include "tensorflow/cc/saved_model/reader.h"
#include "tensorflow/cc/saved_model/signature_constants.h"
#include "tensorflow/cc/saved_model/tag_constants.h"
#include "tensorflow/core/example/example.pb.h"
#include "tensorflow/core/example/feature.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
namespace tensorflow {
namespace {
constexpr char kTestDataPbTxt[] =
"cc/saved_model/testdata/half_plus_two_pbtxt/00000123";
constexpr char kTestDataMainOp[] =
"cc/saved_model/testdata/half_plus_two_main_op/00000123";
constexpr char kTestDataSharded[] =
"cc/saved_model/testdata/half_plus_two/00000123";
constexpr char kTestDataInitOpV2[] =
"cc/saved_model/testdata/half_plus_two_v2/00000123";
constexpr char kTestDataV2DebugInfo[] =
"cc/saved_model/testdata/x_plus_y_v2_debuginfo";
constexpr char kTestFuzzGeneratedNegativeShape[] =
"cc/saved_model/testdata/fuzz_generated/negative_shape";
constexpr char kTestFuzzGeneratedConstWithNoValue[] =
"cc/saved_model/testdata/fuzz_generated/const_with_no_value";
constexpr char kTestFuzzGeneratedBadNodeAttr[] =
"cc/saved_model/testdata/fuzz_generated/bad_node_attr";
constexpr char kTestCyclicModule[] = "cc/saved_model/testdata/CyclicModule";
constexpr char kTestSimpleV1Model[] = "cc/saved_model/testdata/SimpleV1Model";
constexpr char kVarsAndArithmeticObjectGraph[] =
"cc/saved_model/testdata/VarsAndArithmeticObjectGraph";
// This is the value in testdata/VarsAndArithmeticObjectGraph/fingerprint.pb
constexpr char kV2ModuleSavedModelChecksum[] = "15788619162413586750";
class LoaderTest : public ::testing::Test {
protected:
LoaderTest() {}
string MakeSerializedExample(float x) {
tensorflow::Example example;
auto* feature_map = example.mutable_features()->mutable_feature();
(*feature_map)["x"].mutable_float_list()->add_value(x);
return example.SerializeAsString();
}
void ValidateAssets(const string& export_dir,
const SavedModelBundle& bundle) {
const string asset_directory =
io::JoinPath(export_dir, kSavedModelAssetsDirectory);
const string asset_filename = "foo.txt";
const string asset_filepath = io::JoinPath(asset_directory, asset_filename);
TF_EXPECT_OK(Env::Default()->FileExists(asset_filepath));
std::vector<Tensor> path_outputs;
TF_ASSERT_OK(
bundle.session->Run({}, {"filename_tensor:0"}, {}, &path_outputs));
ASSERT_EQ(1, path_outputs.size());
test::ExpectTensorEqual<tstring>(
test::AsTensor<tstring>({"foo.txt"}, TensorShape({})), path_outputs[0]);
}
void CheckSavedModelBundle(const string& export_dir,
const SavedModelBundle& bundle) {
ValidateAssets(export_dir, bundle);
// Retrieve the regression signature from meta graph def.
const auto& signature_def = bundle.GetSignatures().at("regress_x_to_y");
const string input_name = signature_def.inputs().at(kRegressInputs).name();
const string output_name =
signature_def.outputs().at(kRegressOutputs).name();
std::vector<tstring> serialized_examples;
for (float x : {0, 1, 2, 3}) {
serialized_examples.push_back(MakeSerializedExample(x));
}
// Validate the half plus two behavior.
Tensor input =
test::AsTensor<tstring>(serialized_examples, TensorShape({4}));
std::vector<Tensor> outputs;
TF_ASSERT_OK(bundle.session->Run({{input_name, input}}, {output_name}, {},
&outputs));
ASSERT_EQ(outputs.size(), 1);
test::ExpectTensorEqual<float>(
outputs[0],
test::AsTensor<float>({2, 2.5, 3, 3.5}, TensorShape({4, 1})));
}
};
// Test for resource leaks related to TensorFlow session closing requirements
// when loading and unloading large numbers of SavedModelBundles.
// TODO(sukritiramesh): Increase run iterations and move outside of the test
// suite.
TEST_F(LoaderTest, ResourceLeakTest) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
for (int i = 0; i < 100; ++i) {
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
}
TEST_F(LoaderTest, TagMatch) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, ReadMetaGraphFromSavedModel) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
MetaGraphDef actual_metagraph;
TF_ASSERT_OK(ReadMetaGraphDefFromSavedModel(export_dir, {kSavedModelTagServe},
&actual_metagraph));
EXPECT_EQ(actual_metagraph.DebugString(),
bundle.meta_graph_def.DebugString());
}
TEST_F(LoaderTest, RestoreSession) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
SavedModelBundle actual_bundle;
const std::unordered_set<std::string> tags = {kSavedModelTagServe};
TF_ASSERT_OK(ReadMetaGraphDefFromSavedModel(export_dir, tags,
&actual_bundle.meta_graph_def));
TF_ASSERT_OK(LoadMetagraphIntoSession(
session_options, actual_bundle.meta_graph_def, &actual_bundle.session));
TF_ASSERT_OK(RestoreSession(run_options, actual_bundle.meta_graph_def,
export_dir, &actual_bundle.session));
CheckSavedModelBundle(export_dir, actual_bundle);
}
TEST_F(LoaderTest, NoTagMatch) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{"missing-tag"}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(),
"Could not find meta graph def matching supplied tags: { missing-tag }"))
<< st.message();
}
TEST_F(LoaderTest, NoTagMatchMultiple) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st =
LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe, "missing-tag"}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(
st.message(), "Could not find meta graph def matching supplied tags: "))
<< st.message();
}
TEST_F(LoaderTest, SessionCreationFailure) {
SavedModelBundle bundle;
// Use invalid SessionOptions to cause session creation to fail. Default
// options work, so provide an invalid value for the target field.
SessionOptions session_options;
constexpr char kInvalidTarget[] = "invalid target";
session_options.target = kInvalidTarget;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_TRUE(absl::StrContains(st.message(), kInvalidTarget)) << st.message();
}
TEST_F(LoaderTest, PbtxtFormat) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataPbTxt);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, MainOpFormat) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataMainOp);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, InvalidExportPath) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "missing-path");
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
}
TEST_F(LoaderTest, MaybeSavedModelDirectory) {
// Valid SavedModel directory.
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataSharded);
EXPECT_TRUE(MaybeSavedModelDirectory(export_dir));
// Directory that does not exist.
const string missing_export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "missing-path");
EXPECT_FALSE(MaybeSavedModelDirectory(missing_export_dir));
// Directory that exists but is an invalid SavedModel location.
const string invalid_export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model");
EXPECT_FALSE(MaybeSavedModelDirectory(invalid_export_dir));
}
TEST_F(LoaderTest, SavedModelInitOpV2Format) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataInitOpV2);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
CheckSavedModelBundle(export_dir, bundle);
}
TEST_F(LoaderTest, SavedModelV2DebugInfo) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestDataV2DebugInfo);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
// This SavedModel has debug info, so we should have loaded it.
EXPECT_NE(bundle.debug_info.get(), nullptr);
}
TEST_F(LoaderTest, NegativeShapeDimension) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir = io::JoinPath(testing::TensorFlowSrcRoot(),
kTestFuzzGeneratedNegativeShape);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_NE(st.message().find("initializes from a tensor with -1 elements"),
std::string::npos);
}
TEST_F(LoaderTest, ConstNoValue) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir = io::JoinPath(testing::TensorFlowSrcRoot(),
kTestFuzzGeneratedConstWithNoValue);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_NE(st.message().find("constant tensor but no value has been provided"),
std::string::npos);
}
TEST_F(LoaderTest, BadNodeAttr) {
SavedModelBundle bundle;
RunOptions run_options;
SessionOptions session_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestFuzzGeneratedBadNodeAttr);
absl::Status st = LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle);
EXPECT_FALSE(st.ok());
EXPECT_NE(st.message().find("constant tensor but no value has been provided"),
std::string::npos);
}
TEST_F(LoaderTest, UpdateMetricsV2) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string kCCLoadLabel = "cc_load";
const int read_count_v2 = metrics::SavedModelReadCount("2").value();
const int api_count = metrics::SavedModelReadApi(kCCLoadLabel).value();
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestCyclicModule);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
EXPECT_EQ(metrics::SavedModelReadCount("2").value(), read_count_v2 + 1);
EXPECT_EQ(metrics::SavedModelReadApi(kCCLoadLabel).value(), api_count + 1);
}
TEST_F(LoaderTest, UpdateMetricsV1) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string kCCLoadLabel = "cc_load";
const int read_count_v1 = metrics::SavedModelReadCount("1").value();
const int read_count_v2 = metrics::SavedModelReadCount("2").value();
const int api_count = metrics::SavedModelReadApi(kCCLoadLabel).value();
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kTestSimpleV1Model);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
EXPECT_EQ(metrics::SavedModelReadCount("1").value(), read_count_v1 + 1);
EXPECT_EQ(metrics::SavedModelReadCount("2").value(), read_count_v2);
EXPECT_EQ(metrics::SavedModelReadApi(kCCLoadLabel).value(), api_count + 1);
}
TEST_F(LoaderTest, UpdateFingerprintMetrics) {
SavedModelBundle bundle;
SessionOptions session_options;
RunOptions run_options;
const string export_dir =
io::JoinPath(testing::TensorFlowSrcRoot(), kVarsAndArithmeticObjectGraph);
TF_ASSERT_OK(LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle));
EXPECT_EQ(metrics::SavedModelReadPath().value(), export_dir);
EXPECT_EQ(metrics::SavedModelReadFingerprint().value(),
kV2ModuleSavedModelChecksum);
}
} // namespace
} // namespace tensorflow
tensorflow/cc/saved_model/saved_model_fuzz.cc
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2023 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.
==============================================================================*/
#include <memory>
#include "fuzztest/fuzztest.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/cc/saved_model/loader.h"
#include "tensorflow/cc/saved_model/tag_constants.h"
#include "tensorflow/core/protobuf/saved_model.pb.h"
#include "tensorflow/core/public/session_options.h"
namespace tensorflow::fuzzing {
namespace {
void FuzzLoadSavedModel(const SavedModel& model) {
SavedModelBundleLite bundle;
SessionOptions session_options;
RunOptions run_options;
string export_dir = "ram://";
TF_CHECK_OK(tsl::WriteBinaryProto(tensorflow::Env::Default(),
export_dir + kSavedModelFilenamePb, model));
LoadSavedModel(session_options, run_options, export_dir,
{kSavedModelTagServe}, &bundle)
.IgnoreError();
}
FUZZ_TEST(SavedModelFuzz, FuzzLoadSavedModel);
} // namespace
} // namespace tensorflow::fuzzing
tensorflow/cc/saved_model/signature_constants.h
+69
View File
@@ -0,0 +1,69 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_
#define TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_
namespace tensorflow {
/// Key in the signature def map for `default` serving signatures. The default
/// signature is used in inference requests where a specific signature was not
/// specified.
static constexpr char kDefaultServingSignatureDefKey[] = "serving_default";
////////////////////////////////////////////////////////////////////////////////
/// Classification API constants.
/// Classification inputs.
static constexpr char kClassifyInputs[] = "inputs";
/// Classification method name used in a SignatureDef.
static constexpr char kClassifyMethodName[] = "tensorflow/serving/classify";
/// Classification classes output.
static constexpr char kClassifyOutputClasses[] = "classes";
/// Classification scores output.
static constexpr char kClassifyOutputScores[] = "scores";
////////////////////////////////////////////////////////////////////////////////
/// Predict API constants.
/// Predict inputs.
static constexpr char kPredictInputs[] = "inputs";
/// Predict method name used in a SignatureDef.
static constexpr char kPredictMethodName[] = "tensorflow/serving/predict";
/// Predict outputs.
static constexpr char kPredictOutputs[] = "outputs";
////////////////////////////////////////////////////////////////////////////////
/// Regression API constants.
/// Regression inputs.
static constexpr char kRegressInputs[] = "inputs";
/// Regression method name used in a SignatureDef.
static constexpr char kRegressMethodName[] = "tensorflow/serving/regress";
/// Regression outputs.
static constexpr char kRegressOutputs[] = "outputs";
////////////////////////////////////////////////////////////////////////////////
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_SIGNATURE_CONSTANTS_H_
tensorflow/cc/saved_model/tag_constants.h
+35
View File
@@ -0,0 +1,35 @@
/* Copyright 2016 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_
#define TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_
namespace tensorflow {
/// Tag for the `gpu` graph.
constexpr char kSavedModelTagGpu[] = "gpu";
/// Tag for the `tpu` graph.
constexpr char kSavedModelTagTpu[] = "tpu";
/// Tag for the `serving` graph.
constexpr char kSavedModelTagServe[] = "serve";
/// Tag for the `training` graph.
constexpr char kSavedModelTagTrain[] = "train";
} // namespace tensorflow
#endif // TENSORFLOW_CC_SAVED_MODEL_TAG_CONSTANTS_H_
tensorflow/cc/saved_model/test_utils.h
+53
View File
@@ -0,0 +1,53 @@
/* Copyright 2021 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.
==============================================================================*/
#ifndef TENSORFLOW_CC_SAVED_MODEL_TEST_UTILS_H_
#define TENSORFLOW_CC_SAVED_MODEL_TEST_UTILS_H_
#include <ostream>
#include <string>
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow::saved_model {
// TODO(b/229726259) Switch to OSS version after it's available.
// Simple implementation of a proto matcher comparing string representations.
// Only works as ShapeProto's textual representation is deterministic.
class ProtoStringMatcher {
public:
explicit ProtoStringMatcher(const tensorflow::protobuf::Message& expected)
: expected_(expected.DebugString()) {}
template <typename Message>
bool MatchAndExplain(const Message& p,
::testing::MatchResultListener*) const {
return p.DebugString() == expected_;
}
void DescribeTo(::std::ostream* os) const { *os << expected_; }
void DescribeNegationTo(::std::ostream* os) const {
*os << "not equal to expected message: " << expected_;
}
private:
const std::string expected_;
};
inline ::testing::PolymorphicMatcher<ProtoStringMatcher> EqualsProto(
const tensorflow::protobuf::Message& x) {
return ::testing::MakePolymorphicMatcher(ProtoStringMatcher(x));
}
} // namespace tensorflow::saved_model
#endif // TENSORFLOW_CC_SAVED_MODEL_TEST_UTILS_H_
tensorflow/cc/saved_model/testdata/AssetModule/assets/test_asset.txt
Vendored
+1
View File
@@ -0,0 +1 @@
TEST ASSET FILE CONTENTS
tensorflow/cc/saved_model/testdata/AssetModule/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/CyclicModule/debug/saved_model_debug_info.pb
Vendored
+235
View File
@@ -0,0 +1,235 @@
N/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/absl/app.py
m/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/save.py
generate_saved_models.pyb
:AssignVariableOp/MyVariable@__inference__traced_restore_45$
Ö
•
A
ú
«
GW
/AssignVariableOp@__inference__traced_restore_45$
Ö
•
A
ú
«
GO
'Identity@__inference__traced_restore_45$
Ö
•
A
ú
«
GL
$Identity@__inference__traced_save_30$
Ö
•
A
ú
«
GQ
)Identity_1@__inference__traced_restore_45$
Ö
•
A
ú
«
GN
&Identity_1@__inference__traced_save_30$
Ö
•
A
ú
«
GQ
)Identity_2@__inference__traced_restore_45$
Ö
•
A
ú
«
Gj
BMergeV2Checkpoints/checkpoint_prefixes@__inference__traced_save_30$
Ö
•
A
ú
«
GV
.MergeV2Checkpoints@__inference__traced_save_30$
Ö
•
A
ú
«
GK
#NoOp@__inference__traced_restore_45$
Ö
•
A
ú
«
Ga
9RestoreV2/shape_and_slices@__inference__traced_restore_45$
Ö
•
A
ú
«
G]
5RestoreV2/tensor_names@__inference__traced_restore_45$
Ö
•
A
ú
«
GP
(RestoreV2@__inference__traced_restore_45$
Ö
•
A
ú
«
Gc
;RestoreV2_1/shape_and_slices@__inference__traced_restore_45$
Ö
•
A
ú
«
G_
7RestoreV2_1/tensor_names@__inference__traced_restore_45$
Ö
•
A
ú
«
GR
*RestoreV2_1@__inference__traced_restore_45$
Ö
•
A
ú
«
Gi
ASaveV2/MyVariable/Read/ReadVariableOp@__inference__traced_save_30$
Ö
•
A
ú
«
G[
3SaveV2/shape_and_slices@__inference__traced_save_30$
Ö
•
A
ú
«
GW
/SaveV2/tensor_names@__inference__traced_save_30$
Ö
•
A
ú
«
GJ
"SaveV2@__inference__traced_save_30$
Ö
•
A
ú
«
GR
*SaveV2_1/Const@__inference__traced_save_30$
Ö
•
A
ú
«
G]
5SaveV2_1/shape_and_slices@__inference__traced_save_30$
Ö
•
A
ú
«
GY
1SaveV2_1/tensor_names@__inference__traced_save_30$
Ö
•
A
ú
«
GL
$SaveV2_1@__inference__traced_save_30$
Ö
•
A
ú
«
GY
1ShardedFilename/shard@__inference__traced_save_30$
Ö
•
A
ú
«
GS
+ShardedFilename@__inference__traced_save_30$
Ö
•
A
ú
«
G[
3ShardedFilename_1/shard@__inference__traced_save_30$
Ö
•
A
ú
«
GU
-ShardedFilename_1@__inference__traced_save_30$
Ö
•
A
ú
«
GW
/StringJoin/inputs_1@__inference__traced_save_30$
Ö
•
A
ú
«
GN
&StringJoin@__inference__traced_save_30$
Ö
•
A
ú
«
GR
*file_prefix@__inference__traced_restore_45$
Ö
•
A
ú
«
GO
'file_prefix@__inference__traced_save_30$
Ö
•
A
ú
«
GN
&num_shards@__inference__traced_save_30$
Ö
•
A
ú
«
G
tensorflow/cc/saved_model/testdata/CyclicModule/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/CyclicModule/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/CyclicModule/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/OptimizerSlotVariableModule/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/OptimizerSlotVariableModule/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/OptimizerSlotVariableModule/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/SimpleV1Model/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/SimpleV1Model/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/SimpleV1Model/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/StaticHashTableModule/assets/static_hashtable_asset.txt
Vendored
+4
View File
@@ -0,0 +1,4 @@
foo
bar
baz
wombat
tensorflow/cc/saved_model/testdata/StaticHashTableModule/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/StaticHashTableModule/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/StaticHashTableModule/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph/debug/saved_model_debug_info.pb
Vendored
+506
View File
@@ -0,0 +1,506 @@
j/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py
m/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/save.py
€/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/saved_model/signature_serialization.py
i/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/ops/math_ops.py
N/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/absl/app.py
o/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/def_function.py
./generate_saved_models.py
k/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py
q/usr/local/google/home/laurenzo/.local/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.pyg
;AssignVariableOp/variable_x@__inference__traced_restore_101(
Ö
•
0
ú
«
6\
0AssignVariableOp@__inference__traced_restore_101(
Ö
•
0
ú
«
6i
=AssignVariableOp_1/variable_y@__inference__traced_restore_101(
Ö
•
0
ú
«
6^
2AssignVariableOp_1@__inference__traced_restore_101(
Ö
•
0
ú
«
6m
AAssignVariableOp_2/child_variable@__inference__traced_restore_101(
Ö
•
0
ú
«
6^
2AssignVariableOp_2@__inference__traced_restore_101(
Ö
•
0
ú
«
6T
(Identity@__inference__traced_restore_101(
Ö
•
0
ú
«
6P
$Identity@__inference__traced_save_80(
Ö
•
0
ú
«
6J
Identity@__inference_compute_34'
T
þ
0
ú
«
6U
)Identity@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6V
*Identity_1@__inference__traced_restore_101(
Ö
•
0
ú
«
6R
&Identity_1@__inference__traced_save_80(
Ö
•
0
ú
«
6V
*Identity_2@__inference__traced_restore_101(
Ö
•
0
ú
«
6V
*Identity_3@__inference__traced_restore_101(
Ö
•
0
ú
«
6V
*Identity_4@__inference__traced_restore_101(
Ö
•
0
ú
«
6n
BMergeV2Checkpoints/checkpoint_prefixes@__inference__traced_save_80(
Ö
•
0
ú
«
6Z
.MergeV2Checkpoints@__inference__traced_save_80(
Ö
•
0
ú
«
6P
$NoOp@__inference__traced_restore_101(
Ö
•
0
ú
«
6f
:RestoreV2/shape_and_slices@__inference__traced_restore_101(
Ö
•
0
ú
«
6b
6RestoreV2/tensor_names@__inference__traced_restore_101(
Ö
•
0
ú
«
6U
)RestoreV2@__inference__traced_restore_101(
Ö
•
0
ú
«
6h
<RestoreV2_1/shape_and_slices@__inference__traced_restore_101(
Ö
•
0
ú
«
6d
8RestoreV2_1/tensor_names@__inference__traced_restore_101(
Ö
•
0
ú
«
6W
+RestoreV2_1@__inference__traced_restore_101(
Ö
•
0
ú
«
6q
ESaveV2/child_variable/Read/ReadVariableOp@__inference__traced_save_80(
Ö
•
0
ú
«
6_
3SaveV2/shape_and_slices@__inference__traced_save_80(
Ö
•
0
ú
«
6[
/SaveV2/tensor_names@__inference__traced_save_80(
Ö
•
0
ú
«
6m
ASaveV2/variable_x/Read/ReadVariableOp@__inference__traced_save_80(
Ö
•
0
ú
«
6m
ASaveV2/variable_y/Read/ReadVariableOp@__inference__traced_save_80(
Ö
•
0
ú
«
6N
"SaveV2@__inference__traced_save_80(
Ö
•
0
ú
«
6X
,SaveV2_1/Const_1@__inference__traced_save_80(
Ö
•
0
ú
«
6a
5SaveV2_1/shape_and_slices@__inference__traced_save_80(
Ö
•
0
ú
«
6]
1SaveV2_1/tensor_names@__inference__traced_save_80(
Ö
•
0
ú
«
6P
$SaveV2_1@__inference__traced_save_80(
Ö
•
0
ú
«
6]
1ShardedFilename/shard@__inference__traced_save_80(
Ö
•
0
ú
«
6W
+ShardedFilename@__inference__traced_save_80(
Ö
•
0
ú
«
6_
3ShardedFilename_1/shard@__inference__traced_save_80(
Ö
•
0
ú
«
6Y
-ShardedFilename_1@__inference__traced_save_80(
Ö
•
0
ú
«
6k
?StatefulPartitionedCall/args_2@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6k
?StatefulPartitionedCall/args_3@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6k
?StatefulPartitionedCall/args_4@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6k
?StatefulPartitionedCall/args_5@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6d
8StatefulPartitionedCall@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6[
/StringJoin/inputs_1@__inference__traced_save_80(
Ö
•
0
ú
«
6R
&StringJoin@__inference__traced_save_80(
Ö
•
0
ú
«
6C
a@__inference_compute_34'
T
þ
0
ú
«
6N
"a@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6y
2add/ReadVariableOp/resource@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿p
)add/ReadVariableOp@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿c
add@__inference_compute_34E
©
’

Ä
ð
·
Ò
†
ô
¿{
4add_1/ReadVariableOp/resource@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿r
+add_1/ReadVariableOp@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿e
add_1@__inference_compute_34E
©
’

Ä
ð
·
Ò
†
ô
¿g
add_2/y@__inference_compute_34E
©
…

Ä
ð
·
Ò
†
ô
¿e
add_2@__inference_compute_34E
©
…

Ä
ð
·
Ò
†
ô
¿C
b@__inference_compute_34'
T
þ
0
ú
«
6N
"b@__inference_signature_wrapper_45(
ƒ
€
0
ú
«
6W
+file_prefix@__inference__traced_restore_101(
Ö
•
0
ú
«
6S
'file_prefix@__inference__traced_save_80(
Ö
•
0
ú
«
6c
mul@__inference_compute_34E
°
…

Ä
ð
·
Ò
†
ô
¿R
&num_shards@__inference__traced_save_80(
Ö
•
0
ú
«
6}
6truediv/ReadVariableOp/resource@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿t
-truediv/ReadVariableOp@__inference_compute_34C
è
‰

Ä
ð
·
Ò
†
ô
¿g
truediv@__inference_compute_34E
ï
’

Ä
ð
·
Ò
†
ô
¿
tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph/fingerprint.pb
Vendored
+1
View File
@@ -0,0 +1 @@
╬╡ыЛ÷╞°▌ш╓╡пПаСХГ ┤ЫН┘≥к©│O ъДпМ╫²─и╖(╤ПВ╟МШЭш∙2
tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/VarsAndArithmeticObjectGraph/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/bert1/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/bert2/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/chunked_saved_model/chunked_model/saved_model.cpb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/chunked_saved_model/chunked_model/saved_model.pbtxt
Vendored
+2063
View File
File diff suppressed because one or more lines are too long
tensorflow/cc/saved_model/testdata/chunked_saved_model/non_chunked_model/fingerprint.pb
Vendored
+1
View File
@@ -0,0 +1 @@
2(ЕћзмЅ…вА Њ‚¦юЎћхујЋў¶в­ђЪвЋЯЕ«ѓПѕњоУА°ійов®Щ
tensorflow/cc/saved_model/testdata/chunked_saved_model/non_chunked_model/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/chunked_saved_model/non_chunked_model/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/chunked_saved_model/non_chunked_model/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/assets/empty
Vendored
View File
tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/bad_node_attr/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/const_with_no_value/assets/empty
Vendored
View File
tensorflow/cc/saved_model/testdata/fuzz_generated/const_with_no_value/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/const_with_no_value/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/const_with_no_value/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/negative_shape/assets/empty
Vendored
View File
tensorflow/cc/saved_model/testdata/fuzz_generated/negative_shape/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/negative_shape/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/fuzz_generated/negative_shape/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/generate_chunked_models.py
Vendored
+76
View File
@@ -0,0 +1,76 @@
# Copyright 2023 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.
# ==============================================================================
"""Generates GraphDef test data for Merger.
Constructs chunked protos test data containing GraphDefs with lots of nodes and
large nodes for Merger::Read and Merger::Merge.
"""
from collections.abc import Sequence
import os
from absl import app
from absl import flags
import numpy as np
from tensorflow.python.compat import v2_compat
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.lib.io import file_io
from tensorflow.python.module import module
from tensorflow.python.saved_model import loader_impl
from tensorflow.python.saved_model import save
from tensorflow.python.saved_model import save_options
from tensorflow.python.util import compat
from tensorflow.tools.proto_splitter import constants
from tensorflow.tools.proto_splitter.python import saved_model as proto_splitter
SPLITTER_TESTDATA_PATH = flags.DEFINE_string(
"path", None, help="Path to testdata directory.")
def generate_non_chunked_model(non_chunked_dir: str):
root = module.Module()
root.c = constant_op.constant(np.random.random_sample([150, 150]))
constants.debug_set_max_size(80000)
root.get_c = def_function.function(lambda: root.c)
signatures = root.get_c.get_concrete_function()
save.save(root, non_chunked_dir, signatures=signatures,
options=save_options.SaveOptions(experimental_image_format=False))
def generate_chunked_model(non_chunked_dir: str, chunked_dir: str):
saved_model = loader_impl.parse_saved_model(non_chunked_dir)
prefix = file_io.join(compat.as_str(chunked_dir), "saved_model")
file_io.write_string_to_file(f"{prefix}.pbtxt", str(saved_model))
proto_splitter.SavedModelSplitter(saved_model).write(prefix)
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError("Too many command-line arguments.")
main_dir = os.path.join(SPLITTER_TESTDATA_PATH.value, "chunked_saved_model")
non_chunked_dir = os.path.join(main_dir, "non_chunked_model")
generate_non_chunked_model(non_chunked_dir)
chunked_dir = os.path.join(main_dir, "chunked_model")
generate_chunked_model(non_chunked_dir, chunked_dir)
if __name__ == "__main__":
v2_compat.enable_v2_behavior()
app.run(main)
tensorflow/cc/saved_model/testdata/generate_saved_models.py
Vendored
+144
View File
@@ -0,0 +1,144 @@
# Copyright 2019 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.
# ==============================================================================
"""Standalone utility to generate some test saved models."""
import os
from absl import app
from tensorflow.python.client import session as session_lib
from tensorflow.python.compat import v2_compat
from tensorflow.python.eager import def_function
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.module import module
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.saved_model import save_options
from tensorflow.python.saved_model import saved_model
from tensorflow.python.trackable import asset
class VarsAndArithmeticObjectGraph(module.Module):
"""Three vars (one in a sub-module) and compute method."""
def __init__(self):
self.x = variables.Variable(1.0, name="variable_x")
self.y = variables.Variable(2.0, name="variable_y")
self.child = module.Module()
self.child.z = variables.Variable(3.0, name="child_variable")
self.child.c = ops.convert_to_tensor(5.0)
@def_function.function(input_signature=[
tensor_spec.TensorSpec((), dtypes.float32),
tensor_spec.TensorSpec((), dtypes.float32)
])
def compute(self, a, b):
return (a + self.x) * (b + self.y) / (self.child.z) + self.child.c
class ReferencesParent(module.Module):
def __init__(self, parent):
super(ReferencesParent, self).__init__()
self.parent = parent
self.my_variable = variables.Variable(3., name="MyVariable")
# Creates a cyclic object graph.
class CyclicModule(module.Module):
def __init__(self):
super(CyclicModule, self).__init__()
self.child = ReferencesParent(self)
class AssetModule(module.Module):
def __init__(self):
self.asset = asset.Asset(
test.test_src_dir_path("cc/saved_model/testdata/test_asset.txt"))
@def_function.function(input_signature=[])
def read_file(self):
return io_ops.read_file(self.asset)
class StaticHashTableModule(module.Module):
"""A module with an Asset, StaticHashTable, and a lookup function."""
def __init__(self):
self.asset = asset.Asset(
test.test_src_dir_path(
"cc/saved_model/testdata/static_hashtable_asset.txt"))
self.table = lookup_ops.StaticHashTable(
lookup_ops.TextFileInitializer(self.asset, dtypes.string,
lookup_ops.TextFileIndex.WHOLE_LINE,
dtypes.int64,
lookup_ops.TextFileIndex.LINE_NUMBER),
-1)
@def_function.function(
input_signature=[tensor_spec.TensorSpec(shape=None, dtype=dtypes.string)])
def lookup(self, word):
return self.table.lookup(word)
def get_simple_session():
ops.disable_eager_execution()
sess = session_lib.Session()
variables.Variable(1.)
sess.run(variables.global_variables_initializer())
return sess
MODULE_CTORS = {
"VarsAndArithmeticObjectGraph": (VarsAndArithmeticObjectGraph, 2),
"CyclicModule": (CyclicModule, 2),
"AssetModule": (AssetModule, 2),
"StaticHashTableModule": (StaticHashTableModule, 2),
"SimpleV1Model": (get_simple_session, 1),
}
def main(args):
if len(args) != 3:
print("Expected: {export_path} {ModuleName}")
print("Allowed ModuleNames:", MODULE_CTORS.keys())
return 1
_, export_path, module_name = args
module_ctor, version = MODULE_CTORS.get(module_name)
if not module_ctor:
print("Expected ModuleName to be one of:", MODULE_CTORS.keys())
return 2
os.makedirs(export_path)
tf_module = module_ctor()
if version == 2:
options = save_options.SaveOptions(save_debug_info=True)
saved_model.save(tf_module, export_path, options=options)
else:
builder = saved_model.builder.SavedModelBuilder(export_path)
builder.add_meta_graph_and_variables(tf_module, ["serve"])
builder.save()
if __name__ == "__main__":
v2_compat.enable_v2_behavior()
app.run(main)
tensorflow/cc/saved_model/testdata/half_plus_two/00000123/assets/foo.txt
Vendored
+1
View File
@@ -0,0 +1 @@
asset-file-contents
tensorflow/cc/saved_model/testdata/half_plus_two/00000123/saved_model.pb
Vendored Executable
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two/00000123/variables/variables.data-00000-of-00001
Vendored Executable
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two/00000123/variables/variables.index
Vendored Executable
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/assets/foo.txt
Vendored
+1
View File
@@ -0,0 +1 @@
asset-file-contents
tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two_main_op/00000123/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/cc/saved_model/testdata/half_plus_two_pbtxt/00000123/assets/foo.txt
Vendored
+1
View File
@@ -0,0 +1 @@
asset-file-contents

Some files were not shown because too many files have changed in this diff Show More