wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 5 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/c/experimental/saved_model/README.md
+66
View File
@@ -0,0 +1,66 @@
# Tensorflow C SavedModel API
## Overview
These are the new experimental C SavedModel APIs for loading and running
SavedModels in a TF2-idiomatic fashion. See
[RFC 207](https://github.com/tensorflow/community/pull/207) for additional
context.
The directory structure is as follows:
```none
saved_model/
public/
internal/
core/
```
## saved_model/public
`saved_model/public` is intended to house *only the public headers* of the
SavedModel C API.
These headers:
1. declare opaque C types (like `TF_SavedModel`),
2. declare the functions that operate on these types (like `TF_LoadSavedModel`).
Once they leave experimental, these APIs should be considered stable for use
by external clients.
These headers are in a separate directory to make it obvious to clients which
headers they should depend on, and which headers are implementation details.
Separating these public headers by directory also allow future programmatic
checks to ensure that TF public headers only `#include` other public TF headers.
## saved_model/internal
`saved_model/internal` is the "glue" between the C API and the internal C++
implementation.
Its role is to:
1. implement the C API functions declared in `saved_model/public`
2. define the C API types declared in `saved_model/public`
The files fulfilling 1. are named `*.cc` (eg: `concrete_function.cc`), while
the files fulfilling 2. are `*type.h` (eg: `concrete_function_type.h`).
The headers exposing the internal implementation of the opaque C types are only
visible to other implementors of the C API. This is similar to how other
TF C API implementations use `tf_status_internal.h` (to extract the underlying
`tensorflow::Status`). All other targets in this directory are private.
## saved_model/core
`saved_model/core` contains pure C++ "Classes" underlying the C API types
in `saved_model/public/`. These are implementation
details subject to change, and have limited visibility to implementors only.
This is the bottom-most layer of the `C++ -> C -> C++` sandwich.
tensorflow/c/experimental/saved_model/core/BUILD
+310
View File
@@ -0,0 +1,310 @@
# Experimental SavedModel C APIs for TensorFlow. See RFC
# https://github.com/tensorflow/community/pull/207
# Targets in this directory are pure C++ "Classes" underlying the C API types
# under tf/c/experimental/saved_model/public/. They are subject to change and
# have visibility limited to Tensorflow's implementation only.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow/c:__subpackages__",
"//tensorflow/c/experimental/saved_model/internal:__pkg__",
],
licenses = ["notice"],
)
cc_library(
name = "concrete_function",
hdrs = [
"concrete_function.h",
],
deps = [
":function_metadata",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "function_metadata",
hdrs = [
"function_metadata.h",
],
)
cc_library(
name = "saved_model_api",
hdrs = [
"saved_model_api.h",
],
deps = [
":concrete_function",
":signature_def_function",
"//tensorflow/cc/saved_model:bundle_v2",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
],
)
cc_library(
name = "saved_model_utils",
srcs = [
"saved_model_utils.cc",
],
hdrs = [
"saved_model_utils.h",
],
deps = [
":function_metadata",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/experimental/saved_model/core/revived_types:asset",
"//tensorflow/c/experimental/saved_model/core/revived_types:constant",
"//tensorflow/c/experimental/saved_model/core/revived_types:partially_revived_objects",
"//tensorflow/c/experimental/saved_model/core/revived_types:restored_resource_revival_state",
"//tensorflow/c/experimental/saved_model/core/revived_types:tf_concrete_function",
"//tensorflow/c/experimental/saved_model/core/revived_types:tf_concrete_function_revival_state",
"//tensorflow/c/experimental/saved_model/core/revived_types:tf_signature_def_function_revival_state",
"//tensorflow/c/experimental/saved_model/core/revived_types:variable",
"//tensorflow/cc/saved_model:loader_util",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "signature_def_function",
hdrs = [
"signature_def_function.h",
],
deps = [
":signature_def_function_metadata",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "signature_def_function_metadata",
srcs = [
"signature_def_function_metadata.cc",
],
hdrs = [
"signature_def_function_metadata.h",
],
deps = [
":tensor_spec",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
],
)
cc_library(
name = "test_utils",
testonly = True,
srcs = [
"test_utils.cc",
],
hdrs = [
"test_utils.h",
],
deps = [
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "tf_concrete_function_test_protos",
testonly = True,
srcs = ["tf_concrete_function_test_protos.cc"],
hdrs = ["tf_concrete_function_test_protos.h"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "tf_saved_model_api",
srcs = [
"tf_saved_model_api.cc",
],
hdrs = ["tf_saved_model_api.h"],
deps = [
":concrete_function",
":saved_model_api",
":saved_model_utils",
":signature_def_function",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/ops:restore_ops",
"//tensorflow/c/experimental/saved_model/core/revived_types:constant",
"//tensorflow/c/experimental/saved_model/core/revived_types:flat_tensor_function",
"//tensorflow/c/experimental/saved_model/core/revived_types:partially_revived_objects",
"//tensorflow/c/experimental/saved_model/core/revived_types:revived_objects",
"//tensorflow/c/experimental/saved_model/core/revived_types:tensorhandle_convertible",
"//tensorflow/c/experimental/saved_model/core/revived_types:tf_concrete_function",
"//tensorflow/c/experimental/saved_model/core/revived_types:variable",
"//tensorflow/cc/saved_model:bundle_v2",
"//tensorflow/cc/saved_model:constants",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],
)
tf_cc_test(
name = "constant_loading_test",
srcs = [
"constant_loading_test.cc",
],
deps = [
":saved_model_utils",
":test_utils",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/revived_types:constant",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"@com_google_absl//absl/status",
],
)
tf_cc_test(
name = "object_graph_traversal_test",
srcs = [
"object_graph_traversal_test.cc",
],
deps = [
":saved_model_utils",
":test_utils",
":tf_concrete_function_test_protos",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
],
)
tf_cc_test(
name = "saved_variable_loading_test",
srcs = [
"saved_variable_loading_test.cc",
],
deps = [
":saved_model_utils",
":test_utils",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/revived_types:constant",
"//tensorflow/core:all_kernels",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"//tensorflow/core/common_runtime/eager:tensor_handle",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
],
)
tf_cc_test(
name = "signature_flattening_test",
srcs = [
"signature_flattening_test.cc",
],
deps = [
":saved_model_utils",
"//tensorflow/c/experimental/saved_model/core:tf_concrete_function_test_protos",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime/eager:core",
],
)
cc_library(
name = "tensor_spec",
srcs = [
"tensor_spec.cc",
],
hdrs = [
"tensor_spec.h",
],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
],
)
tf_cc_test(
name = "tf_concrete_function_loading_test",
srcs = [
"tf_concrete_function_loading_test.cc",
],
deps = [
":saved_model_utils",
":test_utils",
":tf_concrete_function_test_protos",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/revived_types:constant",
"//tensorflow/c/experimental/saved_model/core/revived_types:tensorhandle_convertible",
"//tensorflow/c/experimental/saved_model/core/revived_types:tf_concrete_function",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"@com_google_absl//absl/status",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tensorflow/c/experimental/saved_model/core/concrete_function.h
+57
View File
@@ -0,0 +1,57 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_CONCRETE_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_CONCRETE_FUNCTION_H_
#include <memory>
#include <vector>
#include "absl/types/span.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
namespace tensorflow {
// ConcreteFunctions correspond to an instance of a tf.function with a known set
// of inputs (either through get_concrete_function) or an input_signature.
// ConcreteFunction attempts to preserve the user-facing semantics of the
// tf.function python API and can take a limited set of types as arguments
// (to be modeled in tensorflow::Value), not just Tensors.
//
// SavedModelAPI's ConcreteFunctions' lifetimes are bound to the SavedModel they
// are loaded from, since they retain pointers to the TensorHandles owned by the
// SavedModel, and the FunctionDef of the SavedModel.
//
// Note(bmzhao): This class is only TEMPORARILY virtual, as a way to unblock
// TFRT integration with TF Serving. Do not add more virtual implementations of
// this class. Eventually we want to remove this virtual base class indirection
// and have only a single implementation.
class ConcreteFunction {
public:
virtual ~ConcreteFunction() = default;
// This method returns the "Call" Op used to execute the function.
virtual absl::Status MakeCallOp(
absl::Span<AbstractTensorHandle* const> inputs,
ImmediateOpPtr* out) const = 0;
virtual const FunctionMetadata& GetFunctionMetadata() const = 0;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_CONCRETE_FUNCTION_H_
tensorflow/c/experimental/saved_model/core/constant_loading_test.cc
+114
View File
@@ -0,0 +1,114 @@
/* 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 <cstdint>
#include <memory>
#include <tuple>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
class ConstantTest : public ::testing::TestWithParam<
std::tuple<DataType, std::vector<int64_t>, bool>> {
public:
ConstantTest()
: device_mgr_(testing::CreateTestingDeviceMgr()),
ctx_(testing::CreateTestingEagerContext(device_mgr_.get())) {}
EagerContext* context() { return ctx_.get(); }
private:
std::unique_ptr<StaticDeviceMgr> device_mgr_;
EagerContextPtr ctx_;
};
// Basic sanity check that roundtripping a Tensor->Tensorproto->Constant
// preserves values.
TEST_P(ConstantTest, CreateConstantSuccessful) {
// Get test parameters
auto& test_params = GetParam();
DataType dtype = std::get<0>(test_params);
TensorShape shape(std::get<1>(test_params));
bool tensorproto_use_tensor_content = std::get<2>(test_params);
// Construct a Tensor with the given dtype + shape
Tensor expected(dtype, shape);
testing::FillNumericTensorBuffer(expected.dtype(), expected.NumElements(),
expected.data(), 42);
// Serialize it to a Tensorproto
TensorProto proto;
if (tensorproto_use_tensor_content) {
expected.AsProtoTensorContent(&proto);
} else {
expected.AsProtoField(&proto);
}
// Revival should succeed w/o errors
std::unique_ptr<Constant> revived;
TF_EXPECT_OK(internal::TensorProtoToConstant(context(), proto, &revived));
// The revived tensorhandle should have the exact same dtype, shape, +
// approx equivalent data to the original.
ImmediateExecutionTensorHandle* handle = revived->handle();
absl::Status status;
AbstractTensorPtr revived_tensor(handle->Resolve(&status));
TF_EXPECT_OK(status) << "Failed to convert tensorhandle to tensor";
EXPECT_EQ(revived_tensor->Type(), expected.dtype());
EXPECT_EQ(revived_tensor->NumElements(), expected.NumElements());
EXPECT_EQ(revived_tensor->NumDims(), expected.dims());
for (int i = 0; i < expected.dims(); ++i) {
EXPECT_EQ(revived_tensor->Dim(i), expected.dim_size(i));
}
testing::CheckBufferDataIsEqual(expected.dtype(), expected.NumElements(),
revived_tensor->Data(), expected.data());
}
// Test against combinations of tensors that are
// 1. Varying dtypes
// 2. Varying shapes
// 3. TensorProto serialized using tensor_content vs repeated type
INSTANTIATE_TEST_SUITE_P(
ConstantIntegerDtypesTest, ConstantTest,
::testing::Combine(
::testing::ValuesIn(testing::DataTypeSetToVector(kDataTypeIsInteger)),
::testing::ValuesIn(testing::InterestingShapes()),
::testing::Values(false, true)));
INSTANTIATE_TEST_SUITE_P(
ConstantFloatingDtypesTest, ConstantTest,
::testing::Combine(::testing::Values(DT_FLOAT, DT_DOUBLE),
::testing::ValuesIn(testing::InterestingShapes()),
::testing::Values(false, true)));
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/function_metadata.h
+27
View File
@@ -0,0 +1,27 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_FUNCTION_METADATA_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_FUNCTION_METADATA_H_
namespace tensorflow {
class FunctionMetadata {
// TODO(bmzhao): Fill in with fields as necessary
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_FUNCTION_METADATA_H_
tensorflow/c/experimental/saved_model/core/object_graph_traversal_test.cc
+370
View File
@@ -0,0 +1,370 @@
/* 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 "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
namespace {
SavedObjectGraph ParseSavedObjectGraph(absl::string_view text_proto) {
SavedObjectGraph value;
CHECK(tensorflow::protobuf::TextFormat::ParseFromString(text_proto, &value));
return value;
}
constexpr absl::string_view kSingleChildFoo = R"(
nodes {
children {
node_id: 1
local_name: "foo"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
)";
constexpr absl::string_view kSingleChildFooWithFuncBar = R"(
nodes {
children {
node_id: 1
local_name: "foo"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
children {
node_id: 2
local_name: "bar"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
function {
concrete_functions: "__inference_my_func_5"
function_spec {
fullargspec {
named_tuple_value {
name: "FullArgSpec"
values {
key: "args"
value {
list_value {
}
}
}
values {
key: "varargs"
value {
none_value {
}
}
}
values {
key: "varkw"
value {
none_value {
}
}
}
values {
key: "defaults"
value {
none_value {
}
}
}
values {
key: "kwonlyargs"
value {
list_value {
}
}
}
values {
key: "kwonlydefaults"
value {
none_value {
}
}
}
values {
key: "annotations"
value {
dict_value {
}
}
}
}
}
input_signature {
tuple_value {
}
}
}
}
}
concrete_functions {
key: "__inference_my_func_5"
value {
canonicalized_input_signature {
tuple_value {
values {
tuple_value {
}
}
values {
dict_value {
}
}
}
}
output_signature {
tensor_spec_value {
shape {
}
dtype: DT_FLOAT
}
}
}
}
)";
// In this graph, foo.baz and bar.wombat should point to the same object.
constexpr absl::string_view kMultiplePathsToChild = R"(
nodes {
children {
node_id: 1
local_name: "foo"
}
children {
node_id: 2
local_name: "bar"
}
children {
node_id: 3
local_name: "signatures"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
children {
node_id: 4
local_name: "baz"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
children {
node_id: 4
local_name: "wombat"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
user_object {
identifier: "signature_map"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
)";
// `foo` has edge `bar`, which has edge `parent` pointing back to `foo`.
constexpr absl::string_view kCycleBetweenParentAndChild = R"(
nodes {
children {
node_id: 1
local_name: "foo"
}
children {
node_id: 2
local_name: "signatures"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
children {
node_id: 3
local_name: "bar"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
user_object {
identifier: "signature_map"
version {
producer: 1
min_consumer: 1
}
}
}
nodes {
children {
node_id: 1
local_name: "parent"
}
user_object {
identifier: "_generic_user_object"
version {
producer: 1
min_consumer: 1
}
}
}
)";
TEST(ObjectGraphTraversalTest, Success) {
SavedObjectGraph object_graph = ParseSavedObjectGraph(kSingleChildFoo);
std::optional<int> node = internal::FindNodeAtPath("foo", object_graph);
ASSERT_TRUE(node.has_value());
EXPECT_EQ(*node, 1);
}
TEST(ObjectGraphTraversalTest, ObjectNotFound) {
SavedObjectGraph object_graph = ParseSavedObjectGraph(kSingleChildFoo);
std::optional<int> node = internal::FindNodeAtPath("bar", object_graph);
EXPECT_FALSE(node.has_value());
}
TEST(ObjectGraphTraversalTest, CaseSensitiveMismatch) {
SavedObjectGraph object_graph = ParseSavedObjectGraph(kSingleChildFoo);
std::optional<int> node = internal::FindNodeAtPath("FOO", object_graph);
EXPECT_FALSE(node.has_value());
}
TEST(ObjectGraphTraversalTest, NestedObjectFound) {
SavedObjectGraph object_graph =
ParseSavedObjectGraph(kSingleChildFooWithFuncBar);
std::optional<int> node = internal::FindNodeAtPath("foo.bar", object_graph);
ASSERT_TRUE(node.has_value());
EXPECT_EQ(*node, 2);
}
TEST(ObjectGraphTraversalTest, MultiplePathsAliasSameObject) {
SavedObjectGraph object_graph = ParseSavedObjectGraph(kMultiplePathsToChild);
std::optional<int> foo_baz_node =
internal::FindNodeAtPath("foo.baz", object_graph);
ASSERT_TRUE(foo_baz_node.has_value());
EXPECT_EQ(*foo_baz_node, 4);
std::optional<int> bar_wombat_node =
internal::FindNodeAtPath("bar.wombat", object_graph);
ASSERT_TRUE(bar_wombat_node.has_value());
EXPECT_EQ(*bar_wombat_node, 4);
EXPECT_EQ(*foo_baz_node, *bar_wombat_node);
}
TEST(ObjectGraphTraversalTest, CyclesAreOK) {
SavedObjectGraph object_graph =
ParseSavedObjectGraph(kCycleBetweenParentAndChild);
std::optional<int> foo = internal::FindNodeAtPath("foo", object_graph);
ASSERT_TRUE(foo.has_value());
EXPECT_EQ(*foo, 1);
std::optional<int> foo_bar =
internal::FindNodeAtPath("foo.bar", object_graph);
ASSERT_TRUE(foo_bar.has_value());
EXPECT_EQ(*foo_bar, 3);
std::optional<int> foo_bar_parent =
internal::FindNodeAtPath("foo.bar.parent", object_graph);
ASSERT_TRUE(foo_bar_parent.has_value());
EXPECT_EQ(*foo_bar_parent, 1);
std::optional<int> foo_bar_parent_bar =
internal::FindNodeAtPath("foo.bar.parent.bar", object_graph);
ASSERT_TRUE(foo_bar_parent_bar.has_value());
EXPECT_EQ(*foo_bar_parent_bar, 3);
EXPECT_EQ(*foo, *foo_bar_parent);
EXPECT_EQ(*foo_bar, *foo_bar_parent_bar);
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/ops/BUILD
+109
View File
@@ -0,0 +1,109 @@
# This package contains written convenience helpers for Eager Operations
# used by SavedModel. Once we autogenerate C++ Eager Op wrappers, we can remove these.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
# Restricting visibility for now
"//tensorflow/c/experimental/saved_model/core:__subpackages__",
],
licenses = ["notice"],
)
cc_library(
name = "restore_ops",
srcs = [
"restore_ops.cc",
],
hdrs = [
"restore_ops.h",
],
deps = [
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/lib/llvm_rtti",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "variable_ops",
srcs = [
"variable_ops.cc",
],
hdrs = [
"variable_ops.h",
],
deps = [
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/lib/llvm_rtti",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
tf_cc_test(
name = "restore_ops_test",
srcs = [
"restore_ops_test.cc",
],
data = [
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
],
deps = [
":restore_ops",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:test_utils",
"//tensorflow/cc/saved_model:constants",
"//tensorflow/core:all_kernels",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
],
)
tf_cc_test(
name = "variable_ops_test",
srcs = [
"variable_ops_test.cc",
],
deps = [
":variable_ops",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:test_utils",
"//tensorflow/core:all_kernels",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/common_runtime:core_cpu_lib",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:core",
"@com_google_absl//absl/status",
],
)
tensorflow/c/experimental/saved_model/core/ops/restore_ops.cc
+118
View File
@@ -0,0 +1,118 @@
/* 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/c/experimental/saved_model/core/ops/restore_ops.h"
#include <cstdint>
#include <string>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/tstring.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace internal {
namespace {
// Creates a scalar string tensorhandle containing a single string `s`
absl::Status CreateStringScalarTensorHandle(ImmediateExecutionContext* ctx,
const std::string& s,
ImmediateTensorHandlePtr* out) {
AbstractTensorPtr tensor(ctx->CreateStringScalar(s));
if (tensor.get() == nullptr) {
return absl::InternalError(
"Failed to create scalar string tensor for checkpoint restore");
}
out->reset(ctx->CreateLocalHandle(tensor.get()));
return absl::Status();
}
// Creates a Rank 1 string tensorhandle containing a single string `s`
absl::Status CreateStringVectorTensorHandle(ImmediateExecutionContext* ctx,
const std::string& s,
ImmediateTensorHandlePtr* out) {
int64_t flat_shape[] = {1};
AbstractTensorPtr tensor(ctx->CreateTensor(DT_STRING, flat_shape));
if (tensor.get() == nullptr) {
return absl::InternalError(
"Failed to create vector string tensor for checkpoint restore");
}
// Use placement new to construct the string, since we don't have
// access to Tensor::flat. This is conceptually equivalent to:
// tensor.flat<tstring>()(0) = s
new (tensor->Data()) tstring(s);
out->reset(ctx->CreateLocalHandle(tensor.get()));
return absl::Status();
}
} // namespace
absl::Status SingleRestore(ImmediateExecutionContext* ctx,
const std::string& prefix,
const std::string& checkpoint_key, DataType dtype,
ImmediateTensorHandlePtr* out) {
// Create the EagerOp
ImmediateOpPtr restore_op(ctx->CreateOperation());
TF_RETURN_IF_ERROR(restore_op->Reset("RestoreV2", "/cpu:0"));
TF_RETURN_IF_ERROR(restore_op->SetAttrTypeList("dtypes", &dtype, 1));
ImmediateTensorHandlePtr prefix_handle;
TF_RETURN_IF_ERROR(
CreateStringScalarTensorHandle(ctx, prefix, &prefix_handle));
ImmediateTensorHandlePtr names_handle;
TF_RETURN_IF_ERROR(
CreateStringVectorTensorHandle(ctx, checkpoint_key, &names_handle));
// Note that empty string is the slice spec used for a non-partitioned
// ResourceVariable:
// https://github.com/tensorflow/tensorflow/blob/06ff30f7ea35098cb68a231a9eb7ff3ff4be4e1e/tensorflow/python/training/saving/saveable_object_util.py#L194
ImmediateTensorHandlePtr shapes_and_slices_handle;
TF_RETURN_IF_ERROR(
CreateStringVectorTensorHandle(ctx, "", &shapes_and_slices_handle));
TF_RETURN_IF_ERROR(restore_op->AddInput(prefix_handle.get()));
TF_RETURN_IF_ERROR(restore_op->AddInput(names_handle.get()));
TF_RETURN_IF_ERROR(restore_op->AddInput(shapes_and_slices_handle.get()));
AbstractTensorHandle* restored_handle = nullptr;
int num_retvals = 1;
TF_RETURN_IF_ERROR(restore_op->Execute(
absl::MakeSpan(&restored_handle, num_retvals), &num_retvals));
AbstractTensorHandlePtr owned_restored_handle(restored_handle);
if (!tensorflow::isa<ImmediateExecutionTensorHandle>(
owned_restored_handle.get())) {
return absl::InternalError("Unexpected tensor handle kind.");
}
out->reset(reinterpret_cast<ImmediateExecutionTensorHandle*>(
owned_restored_handle.release()));
return absl::Status();
}
} // namespace internal
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/ops/restore_ops.h
+42
View File
@@ -0,0 +1,42 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_RESTORE_OPS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_RESTORE_OPS_H_
#include <string>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace internal {
// TODO(bmzhao): Add a function to restore multiple tensors in one call.
// Restores a single non-partioned tensorhandle of dtype `dtype`, using
// checkpoint at `prefix`, with a value stored in `checkpoint_key`.
absl::Status SingleRestore(ImmediateExecutionContext* ctx,
const std::string& prefix,
const std::string& checkpoint_key, DataType dtype,
ImmediateTensorHandlePtr* out);
} // namespace internal
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_RESTORE_OPS_H_
tensorflow/c/experimental/saved_model/core/ops/restore_ops_test.cc
+118
View File
@@ -0,0 +1,118 @@
/* 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/c/experimental/saved_model/core/ops/restore_ops.h"
#include <memory>
#include <string>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
std::string CheckpointPrefix(absl::string_view saved_model_dir) {
return io::JoinPath(testing::TensorFlowSrcRoot(), "cc/saved_model/testdata",
saved_model_dir, kSavedModelVariablesDirectory,
kSavedModelVariablesFilename);
}
class RestoreOpsTest : public ::testing::Test {
public:
RestoreOpsTest()
: device_mgr_(testing::CreateTestingDeviceMgr()),
ctx_(testing::CreateTestingEagerContext(device_mgr_.get())) {}
EagerContext* context() { return ctx_.get(); }
private:
std::unique_ptr<StaticDeviceMgr> device_mgr_;
EagerContextPtr ctx_;
};
// One way of obtaining the checkpointa checkpoint's tensor names is:
// bazel run //tensorflow/python/tools:inspect_checkpoint -- --all_tensors
// --file_name="$CKPT_PREFIX".
// Here are the values for VarsAndArithmeticObjectGraph:
// tensor: child/z/.ATTRIBUTES/VARIABLE_VALUE (float32) []
// 3.0
// tensor: x/.ATTRIBUTES/VARIABLE_VALUE (float32) []
// 1.0
// tensor: y/.ATTRIBUTES/VARIABLE_VALUE (float32) []
// 2.0
TEST_F(RestoreOpsTest, RestoreSuccessful) {
ImmediateTensorHandlePtr x_handle;
TF_EXPECT_OK(internal::SingleRestore(
context(), CheckpointPrefix("VarsAndArithmeticObjectGraph"),
"x/.ATTRIBUTES/VARIABLE_VALUE", DT_FLOAT, &x_handle));
AbstractTensorPtr x = testing::TensorHandleToTensor(x_handle.get());
EXPECT_EQ(x->Type(), DT_FLOAT);
EXPECT_EQ(x->NumElements(), 1);
EXPECT_EQ(x->NumDims(), 0);
EXPECT_FLOAT_EQ(*reinterpret_cast<float*>(x->Data()), 1.0f);
ImmediateTensorHandlePtr y_handle;
TF_EXPECT_OK(internal::SingleRestore(
context(), CheckpointPrefix("VarsAndArithmeticObjectGraph"),
"y/.ATTRIBUTES/VARIABLE_VALUE", DT_FLOAT, &y_handle));
AbstractTensorPtr y = testing::TensorHandleToTensor(y_handle.get());
EXPECT_EQ(y->Type(), DT_FLOAT);
EXPECT_EQ(y->NumElements(), 1);
EXPECT_EQ(y->NumDims(), 0);
EXPECT_FLOAT_EQ(*reinterpret_cast<float*>(y->Data()), 2.0f);
ImmediateTensorHandlePtr z_handle;
TF_EXPECT_OK(internal::SingleRestore(
context(), CheckpointPrefix("VarsAndArithmeticObjectGraph"),
"child/z/.ATTRIBUTES/VARIABLE_VALUE", DT_FLOAT, &z_handle));
AbstractTensorPtr z = testing::TensorHandleToTensor(z_handle.get());
EXPECT_EQ(z->Type(), DT_FLOAT);
EXPECT_EQ(z->NumElements(), 1);
EXPECT_EQ(z->NumDims(), 0);
EXPECT_FLOAT_EQ(*reinterpret_cast<float*>(z->Data()), 3.0f);
}
TEST_F(RestoreOpsTest, BadCheckpointPrefixShouldFail) {
ImmediateTensorHandlePtr x_handle;
absl::Status status = internal::SingleRestore(
context(), CheckpointPrefix("unknown_bad_checkpoint_prefix"),
"x/.ATTRIBUTES/VARIABLE_VALUE", DT_FLOAT, &x_handle);
EXPECT_FALSE(status.ok()) << status.message();
}
TEST_F(RestoreOpsTest, BadCheckpointKeyShouldFail) {
ImmediateTensorHandlePtr x_handle;
absl::Status status = internal::SingleRestore(
context(), CheckpointPrefix("VarsAndArithmeticObjectGraph"),
"bad_checkpoint_key", DT_FLOAT, &x_handle);
EXPECT_FALSE(status.ok()) << status.message();
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/ops/variable_ops.cc
+121
View File
@@ -0,0 +1,121 @@
/* 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/c/experimental/saved_model/core/ops/variable_ops.h"
#include <cstdint>
#include <cstring>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/framework/resource_handle.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/gtl/inlined_vector.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace internal {
absl::Status CreateUninitializedResourceVariable(
ImmediateExecutionContext* ctx, DataType dtype, TensorShape shape,
const char* raw_device_name, ImmediateTensorHandlePtr* handle) {
ImmediateOpPtr varhandle_op(ctx->CreateOperation());
TF_RETURN_IF_ERROR(varhandle_op->Reset("VarHandleOp", raw_device_name));
TF_RETURN_IF_ERROR(varhandle_op->SetAttrType("dtype", dtype));
// Note that if shape is unknown rank, shape.dim_sizes() will be empty, and
// shape.dims() will be -1.
absl::InlinedVector<int64_t, 4UL> dim_sizes = shape.dim_sizes();
TF_RETURN_IF_ERROR(varhandle_op->SetAttrShape(
"shape", reinterpret_cast<const int64_t*>(dim_sizes.data()),
shape.dims()));
TF_RETURN_IF_ERROR(varhandle_op->SetAttrString("container", "", 0));
TF_RETURN_IF_ERROR(
varhandle_op->SetAttrString("shared_name", ResourceHandle::ANONYMOUS_NAME,
strlen(ResourceHandle::ANONYMOUS_NAME)));
AbstractTensorHandle* var_handle = nullptr;
int num_retvals = 1;
TF_RETURN_IF_ERROR(varhandle_op->Execute(
absl::MakeSpan(&var_handle, num_retvals), &num_retvals));
AbstractTensorHandlePtr owned_var_handle(var_handle);
if (!tensorflow::isa<ImmediateExecutionTensorHandle>(
owned_var_handle.get())) {
return absl::InternalError("Unexpected tensor handle kind.");
}
handle->reset(reinterpret_cast<ImmediateExecutionTensorHandle*>(
owned_var_handle.release()));
return absl::Status();
}
absl::Status AssignVariable(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* variable_handle,
DataType dtype,
ImmediateExecutionTensorHandle* value) {
ImmediateOpPtr assign_op(ctx->CreateOperation());
TF_RETURN_IF_ERROR(assign_op->Reset("AssignVariableOp", nullptr));
TF_RETURN_IF_ERROR(assign_op->SetAttrType("dtype", dtype));
TF_RETURN_IF_ERROR(assign_op->AddInput(variable_handle));
TF_RETURN_IF_ERROR(assign_op->AddInput(value));
int num_retvals = 0;
TF_RETURN_IF_ERROR(assign_op->Execute({}, &num_retvals));
return absl::Status();
}
absl::Status ReadVariable(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* variable_handle,
DataType dtype, ImmediateTensorHandlePtr* output) {
ImmediateOpPtr read_op(ctx->CreateOperation());
TF_RETURN_IF_ERROR(read_op->Reset("ReadVariableOp", nullptr));
TF_RETURN_IF_ERROR(read_op->SetAttrType("dtype", dtype));
TF_RETURN_IF_ERROR(read_op->AddInput(variable_handle));
AbstractTensorHandle* value = nullptr;
int num_retvals = 1;
TF_RETURN_IF_ERROR(
read_op->Execute(absl::MakeSpan(&value, num_retvals), &num_retvals));
AbstractTensorHandlePtr owned_value(value);
if (!tensorflow::isa<ImmediateExecutionTensorHandle>(owned_value.get())) {
return absl::InternalError("Unexpected tensor handle kind.");
}
output->reset(
reinterpret_cast<ImmediateExecutionTensorHandle*>(owned_value.release()));
return absl::Status();
}
absl::Status DestroyResource(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* handle) {
ImmediateOpPtr destroy_op(ctx->CreateOperation());
TF_RETURN_IF_ERROR(destroy_op->Reset("DestroyResourceOp", nullptr));
TF_RETURN_IF_ERROR(destroy_op->SetAttrBool("ignore_lookup_error", true));
TF_RETURN_IF_ERROR(destroy_op->AddInput(handle));
int num_retvals = 0;
TF_RETURN_IF_ERROR(destroy_op->Execute({}, &num_retvals));
return absl::Status();
}
} // namespace internal
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/ops/variable_ops.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_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H_
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace internal {
// Executes a VarHandleOp using `ctx`, and fills `handle` with the DT_RESOURCE
// TensorHandle associated with the variable. This is equivalent to creating an
// unitialized TF2 tf.Variable.
// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L1867-L1872
absl::Status CreateUninitializedResourceVariable(
ImmediateExecutionContext* ctx, DataType dtype, TensorShape shape,
const char* raw_device_name, ImmediateTensorHandlePtr* handle);
// Executes an AssignVariableOp using `ctx`, assigning the variable associated
// with `variable_handle` with `value`. `dtype` must be the datatype of the
// underlying variable for `variable_handle`. Note that it is illegal to assign
// a variable to a Tensor with a different dtype than what the variable was
// created with.
absl::Status AssignVariable(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* variable_handle,
DataType dtype,
ImmediateExecutionTensorHandle* value);
// Executes a ReadVariableOp using `ctx`. This reads the underlying variable
// value of `variable_handle` and copies the value to `output`. `dtype` must be
// the dtype of the variable associated with `variable_handle`.
absl::Status ReadVariable(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* variable_handle,
DataType dtype, ImmediateTensorHandlePtr* output);
// Executes DestroyResourceOp on `handle`, using `ctx`. This is equivalent to
// the cleanup that occurs in a tf.Variable's EagerResourceDeleter:
// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/ops/resource_variable_ops.py#L289-L290
absl::Status DestroyResource(ImmediateExecutionContext* ctx,
ImmediateExecutionTensorHandle* handle);
} // namespace internal
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_OPS_VARIABLE_OPS_H_
tensorflow/c/experimental/saved_model/core/ops/variable_ops_test.cc
+99
View File
@@ -0,0 +1,99 @@
/* 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/c/experimental/saved_model/core/ops/variable_ops.h"
#include <memory>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/tensor_interface.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
ImmediateTensorHandlePtr CreateScalarTensorHandle(EagerContext* context,
float value) {
AbstractTensorPtr tensor(context->CreateFloatScalar(value));
ImmediateTensorHandlePtr handle(context->CreateLocalHandle(tensor.get()));
return handle;
}
class VariableOpsTest : public ::testing::Test {
public:
VariableOpsTest()
: device_mgr_(testing::CreateTestingDeviceMgr()),
ctx_(testing::CreateTestingEagerContext(device_mgr_.get())) {}
EagerContext* context() { return ctx_.get(); }
private:
std::unique_ptr<StaticDeviceMgr> device_mgr_;
EagerContextPtr ctx_;
};
// Sanity check for variable creation
TEST_F(VariableOpsTest, CreateVariableSuccessful) {
// Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
ImmediateTensorHandlePtr handle;
TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
context(), DT_FLOAT, {}, nullptr, &handle));
// The created TensorHandle should be a DT_Resource
EXPECT_EQ(handle->DataType(), DT_RESOURCE);
}
// Sanity check for variable destruction
TEST_F(VariableOpsTest, DestroyVariableSuccessful) {
// Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
ImmediateTensorHandlePtr handle;
TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
context(), DT_FLOAT, {}, nullptr, &handle));
// Destroy the variable
TF_EXPECT_OK(internal::DestroyResource(context(), handle.get()));
}
// Sanity check for handle assignment and reading
TEST_F(VariableOpsTest, AssignVariableAndReadSuccessful) {
// Create a DT_Resource TensorHandle that points to a scalar DT_FLOAT tensor
ImmediateTensorHandlePtr variable;
TF_EXPECT_OK(internal::CreateUninitializedResourceVariable(
context(), DT_FLOAT, {}, nullptr, &variable));
// Create a Scalar float TensorHandle with value 42, and assign it to
// the variable.
ImmediateTensorHandlePtr my_value = CreateScalarTensorHandle(context(), 42.0);
TF_EXPECT_OK(internal::AssignVariable(context(), variable.get(), DT_FLOAT,
my_value.get()));
// Read back the value from the variable, and check that it is 42.
ImmediateTensorHandlePtr read_value_handle;
TF_EXPECT_OK(internal::ReadVariable(context(), variable.get(), DT_FLOAT,
&read_value_handle));
absl::Status status;
AbstractTensorPtr read_value(read_value_handle->Resolve(&status));
TF_EXPECT_OK(status);
EXPECT_FLOAT_EQ(42.0, *static_cast<float*>(read_value->Data()));
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/BUILD
+268
View File
@@ -0,0 +1,268 @@
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
# This package contains classes corresponding to Revived SavedObjectGraph types
# used by SavedModel. See https://cs.opensource.google/tensorflow/tensorflow/+/c575e2ba93c442121d98d3f125d83fed1339924d:tensorflow/core/protobuf/saved_object_graph.proto;l=56-62
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
# Restricting visibility for now
"//tensorflow/c/experimental/saved_model/core:__pkg__",
],
licenses = ["notice"],
)
cc_library(
name = "asset",
srcs = [
"asset.cc",
],
hdrs = [
"asset.h",
],
deps = [
":tensorhandle_convertible",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/cc/saved_model:constants",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "constant",
srcs = [
"constant.cc",
],
hdrs = [
"constant.h",
],
deps = [
":tensorhandle_convertible",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "flat_tensor_function",
srcs = [
"flat_tensor_function.cc",
],
hdrs = [
"flat_tensor_function.h",
],
deps = [
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "partially_revived_objects",
srcs = [
"partially_revived_objects.cc",
],
hdrs = [
"partially_revived_objects.h",
],
deps = [
":asset",
":constant",
":restored_resource",
":restored_resource_revival_state",
":revived_objects",
":tf_concrete_function",
":tf_concrete_function_revival_state",
":tf_signature_def_function",
":tf_signature_def_function_revival_state",
":variable",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
"//tensorflow/c/experimental/saved_model/core:tensor_spec",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/lib/llvm_rtti",
"//tensorflow/core/platform:hash",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "restored_resource",
srcs = [
"restored_resource.cc",
],
hdrs = [
"restored_resource.h",
],
deps = [
":tensorhandle_convertible",
":tf_concrete_function",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:lib",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "restored_resource_revival_state",
hdrs = [
"restored_resource_revival_state.h",
],
deps = [
":tf_concrete_function_revival_state",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
],
)
cc_library(
name = "revived_objects",
hdrs = [
"revived_objects.h",
],
deps = [
":asset",
":constant",
":restored_resource",
":tf_concrete_function",
":tf_signature_def_function",
":variable",
"//tensorflow/core:lib",
"@com_google_absl//absl/container:flat_hash_map",
],
)
cc_library(
name = "variable",
srcs = [
"variable.cc",
],
hdrs = [
"variable.h",
],
deps = [
":tensorhandle_convertible",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/ops:variable_ops",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime/eager:context",
"//tensorflow/core/common_runtime/eager:tensor_handle",
"//tensorflow/core/lib/llvm_rtti",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "tensorhandle_convertible",
hdrs = [
"tensorhandle_convertible.h",
],
deps = [
"//tensorflow/c/eager:immediate_execution_tensor_handle",
],
)
cc_library(
name = "tf_concrete_function",
srcs = [
"tf_concrete_function.cc",
],
hdrs = [
"tf_concrete_function.h",
],
deps = [
":flat_tensor_function",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:concrete_function",
"//tensorflow/c/experimental/saved_model/core:function_metadata",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "tf_concrete_function_revival_state",
hdrs = [
"tf_concrete_function_revival_state.h",
],
deps = [
":tf_concrete_function",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "tf_signature_def_function",
srcs = [
"tf_signature_def_function.cc",
],
hdrs = [
"tf_signature_def_function.h",
],
deps = [
":flat_tensor_function",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:signature_def_function",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "tf_signature_def_function_revival_state",
hdrs = [
"tf_signature_def_function_revival_state.h",
],
deps = [
":tf_signature_def_function",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/types:optional",
],
)
tensorflow/c/experimental/saved_model/core/revived_types/asset.cc
+54
View File
@@ -0,0 +1,54 @@
/* 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/c/experimental/saved_model/core/revived_types/asset.h"
#include <memory>
#include <string>
#include <utility>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/cc/saved_model/constants.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
Asset::Asset(ImmediateTensorHandlePtr handle)
: TensorHandleConvertible(std::move(handle)) {}
absl::Status Asset::Create(ImmediateExecutionContext* ctx,
const std::string& saved_model_dir,
const std::string& asset_filename,
std::unique_ptr<Asset>* output) {
std::string abs_path =
io::JoinPath(saved_model_dir, kSavedModelAssetsDirectory, asset_filename);
AbstractTensorPtr tensor(ctx->CreateStringScalar(abs_path));
if (tensor.get() == nullptr) {
return absl::InternalError(absl::StrCat(
"Failed to create scalar string tensor for Asset at path ", abs_path));
}
ImmediateTensorHandlePtr handle(ctx->CreateLocalHandle(tensor.get()));
output->reset(new Asset(std::move(handle)));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/asset.h
+53
View File
@@ -0,0 +1,53 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_ASSET_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_ASSET_H_
#include <memory>
#include <string>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
class Asset : public TensorHandleConvertible {
public:
static absl::Status Create(ImmediateExecutionContext* ctx,
const std::string& saved_model_dir,
const std::string& asset_filename,
std::unique_ptr<Asset>* output);
// Asset is movable, but not copyable.
Asset(Asset&& other) = default;
Asset& operator=(Asset&& other) = default;
~Asset() override = default;
private:
explicit Asset(ImmediateTensorHandlePtr handle);
Asset(const Asset&) = delete;
Asset& operator=(const Asset&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_ASSET_H_
tensorflow/c/experimental/saved_model/core/revived_types/constant.cc
+46
View File
@@ -0,0 +1,46 @@
/* 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/c/experimental/saved_model/core/revived_types/constant.h"
#include <memory>
#include <utility>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
Constant::Constant(ImmediateTensorHandlePtr handle)
: TensorHandleConvertible(std::move(handle)) {}
absl::Status Constant::Create(ImmediateExecutionContext* ctx,
AbstractTensorInterface* tensor,
std::unique_ptr<Constant>* output) {
ImmediateExecutionTensorHandle* handle = ctx->CreateLocalHandle(tensor);
if (handle == nullptr) {
return absl::InternalError("Failed to convert tensor to tensorhandle");
}
output->reset(new Constant(ImmediateTensorHandlePtr(handle)));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/constant.h
+57
View File
@@ -0,0 +1,57 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_CONSTANT_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_CONSTANT_H_
#include <memory>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// This class corresponds to python's tf.constant, which is effectively a
// TensorHandle explicitly initialized to some value.
// For now this doesn't do much beyond wrap Context's CreateLocalHandle method,
// and offer a subclass of TensorHandleConvertible. Note that similar to
// the python's eager mode logic, we bypass calling the "Const" op:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/framework/constant_op.py#L301
class Constant : public TensorHandleConvertible {
public:
static absl::Status Create(ImmediateExecutionContext* ctx,
AbstractTensorInterface* tensor,
std::unique_ptr<Constant>* output);
// RevivedConstant is movable, but not copyable.
Constant(Constant&& other) = default;
Constant& operator=(Constant&& other) = default;
~Constant() override = default;
private:
explicit Constant(ImmediateTensorHandlePtr handle);
Constant(const Constant&) = delete;
Constant& operator=(const Constant&) = delete;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_CONSTANT_H_
tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.cc
+94
View File
@@ -0,0 +1,94 @@
/* 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/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
FlatTensorFunction::FlatTensorFunction(
const std::string& name, std::vector<ImmediateTensorHandlePtr> captures,
ImmediateExecutionContext* ctx)
: name_(name), captures_(std::move(captures)), ctx_(ctx) {}
FlatTensorFunction::~FlatTensorFunction() {
absl::Status status = ctx_->RemoveFunction(name_);
if (!status.ok()) {
LOG(ERROR) << "Failed to remove functiondef " << name_ << ". "
<< status.message();
}
}
absl::Status FlatTensorFunction::Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
ImmediateExecutionContext* ctx, std::unique_ptr<FlatTensorFunction>* out) {
TF_RETURN_IF_ERROR(ctx->AddFunctionDef(*function_def));
std::vector<ImmediateTensorHandlePtr> owned_captures;
owned_captures.reserve(captures.size());
for (ImmediateExecutionTensorHandle* capture : captures) {
capture->Ref();
owned_captures.push_back(ImmediateTensorHandlePtr(capture));
}
out->reset(new FlatTensorFunction(function_def->signature().name(),
std::move(owned_captures), ctx));
return absl::Status();
}
absl::Status FlatTensorFunction::MakeCallOp(
absl::Span<AbstractTensorHandle* const> inputs, ImmediateOpPtr* out) const {
out->reset(ctx_->CreateOperation());
// In eager mode, TF2 python executes functions by constructing an op with
// the name of the functiondef:
// https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L545
// In graph mode, we create a PartitionedCallOp instead:
// https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L573
// TODO(bmzhao): After discussing with Allen, we should execute this via a
// PartitionedCallOp for compatibility with "tooling that assumes functions in
// graphs are PartitionedCallOps".
TF_RETURN_IF_ERROR((*out)->Reset(name_.c_str(), nullptr));
// Adding the user-provided inputs to the function.
TF_RETURN_IF_ERROR((*out)->AddInputList(inputs));
absl::Span<AbstractTensorHandle* const> captures(
reinterpret_cast<AbstractTensorHandle* const*>(captures_.data()),
captures_.size());
// Adding the captures of the function.
TF_RETURN_IF_ERROR((*out)->AddInputList(captures));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h
+90
View File
@@ -0,0 +1,90 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// FlatTensorFunction models a TF2 eager runtime view of a callable function,
// taking + returning flat lists of tensors, including any captures.
// Effectively, it is a thin wrapper around a FunctionDef owned by the
// EagerContext, and any TensorHandle captures associated with the function. The
// MakeCallOp method handles the logic of marshaling captures after the user
// provided inputs automatically.
// Note(bmzhao): This class is mainly intended to house low-level reusable
// function logic between SignatureDefFunction and ConcreteFunction, which
// present higher level interfaces. This type does *not* hold any "function
// metadata".
class FlatTensorFunction {
public:
// Factory for creating a FlatTensorFunction.
//
// Params:
// function_def - The function_def associated with the created
// FlatTensorFunction. FlatTensorFunction will register this
// function_def with `ctx` on creation, and de-register it on
// destruction. function_def must be non-null, but
// otherwise has no lifetime requirements.
// captures - The captured TensorHandles associated with this
// FlatTensorFunction. FlatTensorFunction will participate in
// ownership of the handles (it explicitly increments the refcount
// of each handle, and will decrement them on destruction).
// ctx - A handle to the Tensorflow runtime. This MUST be non-null and
// outlive TFConcreteFunction.
// out - The output FlatTensorFunction.
static absl::Status Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
ImmediateExecutionContext* ctx, std::unique_ptr<FlatTensorFunction>* out);
// This method creates a "Call" Op used to execute the function.
absl::Status MakeCallOp(absl::Span<AbstractTensorHandle* const> inputs,
ImmediateOpPtr* out) const;
~FlatTensorFunction();
private:
FlatTensorFunction(const std::string& name,
std::vector<ImmediateTensorHandlePtr> captures,
ImmediateExecutionContext* ctx);
FlatTensorFunction(const FlatTensorFunction&) = delete;
FlatTensorFunction& operator=(const FlatTensorFunction&) = delete;
// Name of the FunctionDef corresponding to this TFConcreteFunction
std::string name_;
std::vector<ImmediateTensorHandlePtr> captures_;
ImmediateExecutionContext* ctx_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_FLAT_TENSOR_FUNCTION_H_
tensorflow/c/experimental/saved_model/core/revived_types/partially_revived_objects.cc
+544
View File
@@ -0,0 +1,544 @@
/* 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/c/experimental/saved_model/core/revived_types/partially_revived_objects.h"
#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/restored_resource.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/restored_resource_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/core/tensor_spec.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/hash.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace {
using StructuredValueDictEntry =
protobuf::MapPair<std::string, StructuredValue>;
using NamedParamMap =
gtl::FlatMap<absl::string_view, const TensorSpecProto*, StringPieceHasher>;
absl::Status AssertAllCreateResourceFunctionsHaveNoCaptures(
const PartiallyRevivedObjects& objects) {
for (const auto& id_and_resource : objects.restored_resources) {
int node_id = id_and_resource.first;
const RestoredResourceRevivalState& resource = id_and_resource.second;
const TFConcreteFunctionRevivalState* create_resource_fn =
resource.create_resource;
if (create_resource_fn == nullptr) {
return absl::FailedPreconditionError(
absl::StrCat("Resource at node ", node_id,
" did not have a create_resource() function"));
}
const SavedConcreteFunction* saved_create_resource_fn =
create_resource_fn->saved_concrete_func;
if (!saved_create_resource_fn->bound_inputs().empty()) {
// TODO(b/124045874): Support loading resource functions via a top sort
return absl::UnimplementedError(
"Create Resource functions with captures are currently unsupported.");
}
}
return absl::Status();
}
// Retrieves the TensorHandle associated with `node_id` from `obj_graph`, and
// set `*handle` to point to it.
absl::Status TensorHandleFromNode(int node_id,
const SavedObjectGraph& obj_graph,
const PartiallyRevivedObjects& objects,
ImmediateExecutionTensorHandle** handle) {
const SavedObject& node = obj_graph.nodes(node_id);
SavedObject::KindCase kind = node.kind_case();
switch (kind) {
case SavedObject::kVariable: {
const auto& variables_iter = objects.variables.find(node_id);
if (variables_iter == objects.variables.end()) {
return absl::FailedPreconditionError(absl::StrCat(
"Tried to convert node id ", node_id,
" of type variable to tensor but the variable wasn't initialized"));
}
*handle = variables_iter->second->handle();
return absl::Status();
}
case SavedObject::kConstant: {
const auto& constants_iter = objects.constants.find(node_id);
if (constants_iter == objects.constants.end()) {
return absl::FailedPreconditionError(
absl::StrCat("Tried to convert node id ", node_id,
" of type constant to tensor but the "
"constant wasn't initialized"));
}
*handle = constants_iter->second->handle();
return absl::Status();
}
case SavedObject::kAsset: {
const auto& assets_iter = objects.assets.find(node_id);
if (assets_iter == objects.assets.end()) {
return absl::FailedPreconditionError(absl::StrCat(
"Tried to convert node id ", node_id,
" of type asset to tensor but the asset wasn't initialized"));
}
*handle = assets_iter->second->handle();
return absl::Status();
}
case SavedObject::kResource: {
const auto& resource_iter = objects.restored_resources.find(node_id);
if (resource_iter == objects.restored_resources.end()) {
return absl::FailedPreconditionError(absl::StrCat(
"Tried to convert node id ", node_id,
" of type Resource to tensor but the Resource wasn't initialized"));
}
const RestoredResourceRevivalState& resource = resource_iter->second;
if (resource.resource_handle == nullptr) {
return absl::FailedPreconditionError(absl::StrCat(
"Resource with node id ", node_id,
" should have its resource_handle created, but was nullptr."));
}
*handle = resource.resource_handle.get();
return absl::Status();
}
default: {
return absl::FailedPreconditionError(absl::StrCat(
"Only objects of type variable, constant, asset, and resources have "
"capturable tensorhandles. Encountered object of kind ",
node.kind_case(), " at node id: ", node_id));
}
}
}
std::vector<SignatureDefParam> SignatureDefParamsFromNamedParamMap(
const NamedParamMap& params) {
// The underlying functiondef associated with the SignatureDef has
// nest.flattened inputs and outputs, which are sorted by string key.
std::vector<SignatureDefParam> result;
result.reserve(params.size());
for (const auto& named_param : params) {
result.push_back(SignatureDefParam(std::string(named_param.first),
TensorSpec(*named_param.second)));
}
std::sort(result.begin(), result.end(),
[](const SignatureDefParam& x, const SignatureDefParam& y) {
return x.name() < y.name();
});
return result;
}
// SignatureDefArgsFromInputs takes the "canonicalized_input_signature"
// field of a SavedConcreteFunction, ensures it conforms to the structure of
// tuple(tuple(), dict<string,TensorSpec>()), and "returns" a list of
// SignatureDefParams of the SignatureDefFunction's arguments.
absl::Status SignatureDefArgsFromInputs(
const StructuredValue& canonicalized_input_signature,
std::vector<SignatureDefParam>* out) {
// Note(bmzhao): canonicalized_input_signature should be a tuple of
// (args, kwargs), where args is an empty tuple, and kwargs is a dictionary of
// string keys to TensorSpecs.
if (!canonicalized_input_signature.has_tuple_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's canonicalized_input_signature should be "
"of form tuple(tuple(), dict()), but was instead: \n",
canonicalized_input_signature.DebugString()));
}
const TupleValue& args_kwargs_tuple =
canonicalized_input_signature.tuple_value();
if (args_kwargs_tuple.values_size() != 2) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's canonicalized_input_signature should be "
"a tuple of two elements (args, kwargs), but was instead: \n",
args_kwargs_tuple.DebugString()));
}
const StructuredValue& args = args_kwargs_tuple.values(0);
if (!args.has_tuple_value() || !args.tuple_value().values().empty()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's canonicalized_input_signature's args"
"should be an empty tuple, but instead got: \n",
args.DebugString()));
}
const StructuredValue& kwargs = args_kwargs_tuple.values(1);
if (!kwargs.has_dict_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's canonicalized_input_signature's kwargs"
"should be a dictionary, but instead got: \n",
kwargs.DebugString()));
}
const DictValue& kwargs_dict = kwargs.dict_value();
NamedParamMap result;
result.reserve(kwargs_dict.fields_size());
for (const auto& key_value : kwargs_dict.fields()) {
const std::string& key = key_value.first;
const StructuredValue& value = key_value.second;
if (!value.has_tensor_spec_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's canonicalized_input_signature's kwargs"
"dictionary contained a non-tensorspec value for key-value pair: \n",
"Key: ", key, "Value: \n", value.DebugString()));
}
result[key] = &value.tensor_spec_value();
}
*out = SignatureDefParamsFromNamedParamMap(result);
return absl::Status();
}
// SignatureDefReturnsFromOutputs takes the "output_signature" field of a
// SavedConcreteFunction, ensures it conforms to the structure of
// dict<string,TensorSpec>(), and "returns" a list of SignatureDefParams of the
// SignatureDefFunction's returns.
absl::Status SignatureDefReturnsFromOutputs(
const StructuredValue& output_signature,
std::vector<SignatureDefParam>* out) {
if (!output_signature.has_dict_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's output_signature must be a dictionary, but "
"instead got: ",
output_signature.DebugString()));
}
const DictValue& output_dict = output_signature.dict_value();
NamedParamMap result;
result.reserve(output_dict.fields_size());
for (const auto& key_value : output_dict.fields()) {
const std::string& key = key_value.first;
const StructuredValue& value = key_value.second;
if (!value.has_tensor_spec_value()) {
return absl::FailedPreconditionError(absl::StrCat(
"SignatureDefFunction's output_signature dictionary contained a "
"non-tensorspec value for key-value pair: \n",
"Key: ", key, "Value: \n", value.DebugString()));
}
result[key] = &value.tensor_spec_value();
}
*out = SignatureDefParamsFromNamedParamMap(result);
return absl::Status();
}
// The implementation takes advantage of the fact that SignatureDefFunction's
// "traced" Signature wrapper function always has inputs/outputs of dictionaries
// https://github.com/tensorflow/tensorflow/blob/53cdd5e87c423b195f33775753273286fd5a1a65/tensorflow/python/saved_model/signature_serialization.py#L119-L126
// https://github.com/tensorflow/tensorflow/blob/53cdd5e87c423b195f33775753273286fd5a1a65/tensorflow/python/saved_model/signature_serialization.py#L153-L178
// Additionally, we take advantage of the fact that the SignatureDefFunction's
// associated functiondef has lexicographically ordered inputs/outputs due to
// nest.flatten.
absl::Status LoadSignatureDefFunctionMetadata(
const SavedConcreteFunction& saved_concrete_function,
SignatureDefFunctionMetadata* out) {
std::vector<SignatureDefParam> args;
TF_RETURN_IF_ERROR(SignatureDefArgsFromInputs(
saved_concrete_function.canonicalized_input_signature(), &args));
std::vector<SignatureDefParam> rets;
TF_RETURN_IF_ERROR(SignatureDefReturnsFromOutputs(
saved_concrete_function.output_signature(), &rets));
*out = SignatureDefFunctionMetadata(std::move(args), std::move(rets));
return absl::Status();
}
// This function finds the necessary captures, then forwards to the builder
// method
absl::Status CreateConcreteFunction(
ImmediateExecutionContext* ctx,
const TFConcreteFunctionRevivalState& builder,
const SavedObjectGraph& obj_graph, const PartiallyRevivedObjects& objects,
std::unique_ptr<TFConcreteFunction>* out) {
const auto& capture_node_ids = builder.saved_concrete_func->bound_inputs();
std::vector<ImmediateExecutionTensorHandle*> captures;
captures.reserve(capture_node_ids.size());
for (int capture_node_id : capture_node_ids) {
ImmediateExecutionTensorHandle* capture_handle;
TF_RETURN_IF_ERROR(TensorHandleFromNode(capture_node_id, obj_graph, objects,
&capture_handle));
captures.push_back(capture_handle);
}
// TODO(bmzhao): Create Metadata here
return TFConcreteFunction::Create(/*function_def=*/builder.fdef,
/*captures=*/std::move(captures),
/*metadata=*/{},
/*ctx=*/ctx,
/*out=*/out);
}
absl::Status CreateSignatureDefFunction(
ImmediateExecutionContext* ctx,
const TFSignatureDefFunctionRevivalState& builder,
const SavedObjectGraph& obj_graph, const PartiallyRevivedObjects& objects,
std::unique_ptr<TFSignatureDefFunction>* out) {
const auto& capture_node_ids = builder.saved_concrete_func->bound_inputs();
std::vector<ImmediateExecutionTensorHandle*> captures;
captures.reserve(capture_node_ids.size());
for (int capture_node_id : capture_node_ids) {
ImmediateExecutionTensorHandle* capture_handle;
TF_RETURN_IF_ERROR(TensorHandleFromNode(capture_node_id, obj_graph, objects,
&capture_handle));
captures.push_back(capture_handle);
}
SignatureDefFunctionMetadata metadata;
TF_RETURN_IF_ERROR(LoadSignatureDefFunctionMetadata(
*builder.saved_concrete_func, &metadata));
return TFSignatureDefFunction::Create(/*function_def=*/builder.fdef,
/*captures=*/std::move(captures),
/*metadata=*/std::move(metadata),
/*ctx=*/ctx,
/*out=*/out);
}
absl::Status InitializeCreateResourceFunctions(
ImmediateExecutionContext* ctx, const SavedObjectGraph& obj_graph,
const PartiallyRevivedObjects& objects, RevivedObjects* revived) {
for (const auto& id_and_resource : objects.restored_resources) {
const RestoredResourceRevivalState& resource = id_and_resource.second;
const TFConcreteFunctionRevivalState* create_resource_fn =
resource.create_resource;
const SavedConcreteFunction* saved_create_resource_fn =
create_resource_fn->saved_concrete_func;
if (!saved_create_resource_fn->bound_inputs().empty()) {
// TODO(b/124045874): Load resource functions via a topological sort
return absl::UnimplementedError(
"Create Resource functions with captures are currently unsupported.");
}
std::unique_ptr<TFConcreteFunction> out;
TF_RETURN_IF_ERROR(CreateConcreteFunction(ctx, *create_resource_fn,
obj_graph, objects, &out));
revived->concrete_functions.Insert(std::move(out),
create_resource_fn->node_id);
}
return absl::Status();
}
absl::Status InitializeAllFunctions(ImmediateExecutionContext* ctx,
const SavedObjectGraph& obj_graph,
const PartiallyRevivedObjects& objects,
RevivedObjects* revived) {
gtl::FlatMap<int, std::unique_ptr<TFSignatureDefFunction>>*
destination_sig_map = &revived->signature_def_functions;
for (const auto& id_and_func : objects.concrete_functions) {
int node_id = id_and_func.first;
const TFConcreteFunctionRevivalState& func = id_and_func.second;
if (revived->concrete_functions.Find(node_id)) {
// The function has already been initialized in the destination_map,
// so we can skip this node. This can occur because we initialize
// CreateResource functions before calling this function.
continue;
}
std::unique_ptr<TFConcreteFunction> out;
TF_RETURN_IF_ERROR(
CreateConcreteFunction(ctx, func, obj_graph, objects, &out));
revived->concrete_functions.Insert(std::move(out), node_id);
}
for (const auto& id_and_func : objects.signature_def_functions) {
int node_id = id_and_func.first;
const TFSignatureDefFunctionRevivalState& func = id_and_func.second;
if (destination_sig_map->find(node_id) != destination_sig_map->end()) {
continue;
}
std::unique_ptr<TFSignatureDefFunction> out;
TF_RETURN_IF_ERROR(
CreateSignatureDefFunction(ctx, func, obj_graph, objects, &out));
(*destination_sig_map)[node_id] = std::move(out);
}
return absl::Status();
}
absl::Status CreateAllResourceHandles(ImmediateExecutionContext* ctx,
const SavedObjectGraph& obj_graph,
PartiallyRevivedObjects* objects,
RevivedObjects* revived) {
for (auto& id_and_resource : objects->restored_resources) {
RestoredResourceRevivalState& resource = id_and_resource.second;
int create_resource_fn_node = resource.create_resource->node_id;
const TFConcreteFunction* create_resource_fn =
revived->concrete_functions.Find(create_resource_fn_node);
if (create_resource_fn == nullptr) {
return absl::FailedPreconditionError(
absl::StrCat("ConcreteFunction at node ", create_resource_fn_node,
" should have been initialized prior to being called."));
}
ImmediateOpPtr function_op;
TF_RETURN_IF_ERROR(create_resource_fn->MakeCallOp({}, &function_op));
TF_RETURN_IF_ERROR(function_op->SetDeviceName(resource.device.c_str()));
AbstractTensorHandle* resource_handle = nullptr;
int num_retvals = 1;
TF_RETURN_IF_ERROR(function_op->Execute(
absl::MakeSpan(&resource_handle, num_retvals), &num_retvals));
AbstractTensorHandlePtr owned_resource_handle(resource_handle);
if (!tensorflow::isa<ImmediateExecutionTensorHandle>(
owned_resource_handle.get())) {
return absl::InternalError("Unexpected tensor handle kind.");
}
ImmediateTensorHandlePtr result(
reinterpret_cast<ImmediateExecutionTensorHandle*>(
owned_resource_handle.release()));
resource.resource_handle = std::move(result);
}
return absl::Status();
}
absl::Status BuildResources(ImmediateExecutionContext* ctx,
const SavedObjectGraph& obj_graph,
PartiallyRevivedObjects* objects,
RevivedObjects* revived) {
for (auto& id_and_resource : objects->restored_resources) {
int node_id = id_and_resource.first;
RestoredResourceRevivalState& resource_revival_state =
id_and_resource.second;
TFConcreteFunction* create_resource = nullptr;
// Check all the functions associated with the resource have already been
// initialized in `revived`
if (resource_revival_state.create_resource != nullptr) {
create_resource = revived->concrete_functions.Find(
resource_revival_state.create_resource->node_id);
if (create_resource == nullptr) {
return absl::FailedPreconditionError(absl::StrCat(
"'create_resource' function with node id ",
resource_revival_state.create_resource->node_id, " not found"));
}
}
TFConcreteFunction* initialize = nullptr;
if (resource_revival_state.initialize != nullptr) {
initialize = revived->concrete_functions.Find(
resource_revival_state.initialize->node_id);
if (initialize == nullptr) {
return absl::FailedPreconditionError(absl::StrCat(
"'initialize' function with node id ",
resource_revival_state.initialize->node_id, " not found"));
}
}
TFConcreteFunction* destroy_resource = nullptr;
if (resource_revival_state.destroy_resource != nullptr) {
destroy_resource = revived->concrete_functions.Find(
resource_revival_state.destroy_resource->node_id);
if (destroy_resource == nullptr) {
return absl::FailedPreconditionError(absl::StrCat(
"'destroy_resource' function with node id ",
resource_revival_state.destroy_resource->node_id, " not found"));
}
}
if (resource_revival_state.resource_handle == nullptr) {
return absl::FailedPreconditionError(
absl::StrCat("Resource at node id ", node_id,
" does not have a resource handle."));
}
revived->restored_resources.emplace(
node_id, RestoredResource(
/*device=*/resource_revival_state.device,
/*create_resource=*/create_resource,
/*initialize=*/initialize,
/*destroy_resource=*/destroy_resource,
/*resource_handle=*/
std::move(resource_revival_state.resource_handle)));
}
return absl::Status();
}
} // namespace
absl::Status PartiallyRevivedObjects::Build(ImmediateExecutionContext* ctx,
const SavedObjectGraph& obj_graph,
RevivedObjects* revived) {
// Step 1: We would like to initialize all functions; this requires setting up
// their captured tensorhandles, which may come from variables, assets,
// constants, or resources. The first three are trivial; However,
// tensorhandles that correspond to resources must be created by invoking
// their "create_resource" function.
// https://github.com/tensorflow/tensorflow/blob/f19c6efb4a8ba60e2492eedc98ef5375abb39dc7/tensorflow/python/saved_model/load.py#L240
// https://github.com/tensorflow/tensorflow/blob/f19c6efb4a8ba60e2492eedc98ef5375abb39dc7/tensorflow/python/training/tracking/tracking.py#L233
// For now, we assert that all create_resource functions must have no
// captures. This aligns with the current behavior in python.
// https://github.com/tensorflow/tensorflow/blob/50eac986bf7a0ad12594e080f083181f277e0b49/tensorflow/python/saved_model/load.py#L152-L155
// TODO(bmzhao): We should do a topological sort instead.
// 1a. Make sure all CreateResource functions have no captures.
TF_RETURN_IF_ERROR(AssertAllCreateResourceFunctionsHaveNoCaptures(*this));
// 1b. Initialize all CreateResource functions, storing them in `revived`
TF_RETURN_IF_ERROR(
InitializeCreateResourceFunctions(ctx, obj_graph, *this, revived));
// 1c. Invoke all "CreateResource" functions and store their ResourceHandles
// https://github.com/tensorflow/tensorflow/blob/3b6b41b68a95dc70c26dc816b29d359bfb88c116/tensorflow/python/training/tracking/tracking.py#L241-L247
// in *this->resources.
// TODO(bmzhao): Maybe store them separately, not in *this?
TF_RETURN_IF_ERROR(CreateAllResourceHandles(ctx, obj_graph, this, revived));
// 2. Initialize all the rest of the functions
TF_RETURN_IF_ERROR(InitializeAllFunctions(ctx, obj_graph, *this, revived));
// 3a. Move over all non-function, non-resource objects
revived->variables = std::move(variables);
revived->assets = std::move(assets);
revived->constants = std::move(constants);
revived->signatures_map = std::move(signatures_map);
// 3b. Move over resources.
TF_RETURN_IF_ERROR(BuildResources(ctx, obj_graph, this, revived));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/partially_revived_objects.h
+65
View File
@@ -0,0 +1,65 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_PARTIALLY_REVIVED_OBJECTS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_PARTIALLY_REVIVED_OBJECTS_H_
#include <memory>
#include <string>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/asset.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/restored_resource_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// Container for objects during the revival step in SavedModel's loading.
// Notably, resources and functions can be in a state where they reference
// other resources/functions that have not been constructed yet. We collect
// *all* objects in a partially valid state here, then properly initialize
// resources and functions. Implementation-wise, PartiallyRevivedObjects
// contains maps keyed by the node number of the SavedObjectGraph, and map to an
// object of the corresponding type. So, if node 2 in the object graph is a
// variable, PartiallyRevivedObjects.variables[2] exists, and corresponds to a
// tensorflow::Variable object. The only exception to this is the
// "signatures_map", which is keyed by the "signature" key
// (https://github.com/tensorflow/tensorflow/blob/372918decee7f558b3c194b04f77c20dcc679a31/tensorflow/core/protobuf/meta_graph.proto#L89),
// and maps to the SignatureDefFunction node in the SavedObjectGraph.
struct PartiallyRevivedObjects {
gtl::FlatMap<int, std::unique_ptr<Variable>> variables;
gtl::FlatMap<int, std::unique_ptr<Asset>> assets;
gtl::FlatMap<int, std::unique_ptr<Constant>> constants;
gtl::FlatMap<int, TFConcreteFunctionRevivalState> concrete_functions;
gtl::FlatMap<int, TFSignatureDefFunctionRevivalState> signature_def_functions;
gtl::FlatMap<int, RestoredResourceRevivalState> restored_resources;
gtl::FlatMap<std::string, int> signatures_map;
absl::Status Build(ImmediateExecutionContext* ctx,
const SavedObjectGraph& obj_graph,
RevivedObjects* revived);
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_PARTIALLY_REVIVED_OBJECTS_H_
tensorflow/c/experimental/saved_model/core/revived_types/restored_resource.cc
+84
View File
@@ -0,0 +1,84 @@
/* 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/c/experimental/saved_model/core/revived_types/restored_resource.h"
#include <string>
#include <utility>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
namespace {
absl::Status ExecuteNoArgDummyReturnFunction(TFConcreteFunction* func) {
ImmediateOpPtr function_op;
TF_RETURN_IF_ERROR(func->MakeCallOp({}, &function_op));
AbstractTensorHandle* dummy_output = nullptr;
int num_retvals = 1;
TF_RETURN_IF_ERROR(function_op->Execute(
absl::MakeSpan(&dummy_output, num_retvals), &num_retvals));
AbstractTensorHandlePtr owned_dummy_output(dummy_output);
return absl::Status();
}
} // namespace
RestoredResource::RestoredResource(const std::string& device,
TFConcreteFunction* create_resource,
TFConcreteFunction* initialize,
TFConcreteFunction* destroy_resource,
ImmediateTensorHandlePtr resource_handle)
: TensorHandleConvertible(std::move(resource_handle)),
device_(device),
create_resource_(create_resource),
initialize_(initialize),
destroy_resource_(destroy_resource) {}
absl::Status RestoredResource::Initialize() const {
return ExecuteNoArgDummyReturnFunction(initialize_);
}
RestoredResource::~RestoredResource() {
// Note(bmzhao): SavedModels saved before
// https://github.com/tensorflow/tensorflow/commit/3c806101f57768e479f8646e7518bbdff1632ca3
// did not have their destroy_resource function saved, meaning they will
// leak resources.
//
// Check that handle is null before calling destroy_resource function in case
// destructor is invoked unintentionally.
if (destroy_resource_ != nullptr && handle() != nullptr) {
absl::Status status = ExecuteNoArgDummyReturnFunction(destroy_resource_);
if (!status.ok()) {
LOG(WARNING)
<< "Failed executing destroy_resource function for RestoredResource: "
<< status.message();
}
}
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/restored_resource.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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_
#include <memory>
#include <string>
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// RestoredResource represents a TF2 "Resource" object loaded from a savedmodel,
// analogous to the Python _RestoredResource object:
// https://github.com/tensorflow/tensorflow/blob/fda326e542ca67534e8411edb180e8760a4828b7/tensorflow/python/saved_model/load.py#L481
// TF2 resource objects typically extend TrackableResource:
// https://github.com/tensorflow/tensorflow/blob/fda326e542ca67534e8411edb180e8760a4828b7/tensorflow/python/training/tracking/tracking.py#L285
// and are expected to implement "_create_resource", "_initialize", and
// "_destroy_resource" functions:
// https://github.com/tensorflow/tensorflow/blob/139ba9c5284799beafdd1d7f895127cf00e7c48f/tensorflow/python/training/tracking/tracking.py#L262-L281
class RestoredResource : TensorHandleConvertible {
public:
// Note(bmzhao): RestoredResource stores non-owning pointers to its associated
// functions because SavedModel internally owns all functions and objects in
// the RevivedObjects struct (which owns all functions). One alternative would
// be to have RevivedObjects store shared_ptr<TFConcreteFunction> instead, and
// change RestoredResource's constructor take shared_ptr<TFConcreteFunction>.
// To keep things simple, I've stuck to raw pointers for now.
//
// Params:
// device - The device string associated with the SavedResource
// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/saved_object_graph.proto#L182
// Conceptually, this is the same device used in CapturableResource:
// https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/python/training/tracking/tracking.py#L222-L225
// Implementation-wise, it is device used when invoking the
// create_resource function to produce the resource_handle
// associated with the object:
// https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/python/training/tracking/tracking.py#L246-L247
// create_resource - Non owning pointer to the create_resource function
// associated with this object. Must be NON-NULL.
// initialize - Non owning pointer to the initialize function associated with
// this object. Must be NON-NULL.
// destroy_resource - Non owning pointer to the destroy_resource function
// associated with this object. Ideally this should be
// NON-NULL, but in order to support models saved prior to
// https://github.com/tensorflow/tensorflow/commit/3c806101f57768e479f8646e7518bbdff1632ca3
// we allow null here. This will, however, leak resources.
RestoredResource(const std::string& device,
TFConcreteFunction* create_resource,
TFConcreteFunction* initialize,
TFConcreteFunction* destroy_resource,
ImmediateTensorHandlePtr resource_handle);
absl::Status Initialize() const;
// RestoredResource is movable, but not copyable.
RestoredResource(RestoredResource&& other) = default;
RestoredResource& operator=(RestoredResource&& other) = default;
~RestoredResource() override;
private:
std::string device_;
TFConcreteFunction* create_resource_;
TFConcreteFunction* initialize_;
TFConcreteFunction* destroy_resource_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_H_
tensorflow/c/experimental/saved_model/core/revived_types/restored_resource_revival_state.h
+38
View File
@@ -0,0 +1,38 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_REVIVAL_STATE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_REVIVAL_STATE_H_
#include <string>
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function_revival_state.h"
namespace tensorflow {
// All "Resources" should have these 3 saved functions:
// https://github.com/tensorflow/tensorflow/blob/86dc281333d7d277ddc1882f2bca4b17e7ec40e5/tensorflow/python/training/tracking/tracking.py#L277-L281
struct RestoredResourceRevivalState {
std::string device;
TFConcreteFunctionRevivalState* create_resource = nullptr;
TFConcreteFunctionRevivalState* initialize = nullptr;
TFConcreteFunctionRevivalState* destroy_resource = nullptr;
ImmediateTensorHandlePtr resource_handle = nullptr;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_RESTORED_RESOURCE_REVIVAL_STATE_H_
tensorflow/c/experimental/saved_model/core/revived_types/revived_objects.h
+92
View File
@@ -0,0 +1,92 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_REVIVED_OBJECTS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_REVIVED_OBJECTS_H_
#include <memory>
#include <unordered_map>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/asset.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/restored_resource.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
namespace tensorflow {
// A container for revived saved model objects.
//
// Most of the objects will be revived from nodes in the object graph, and for
// those objects this container provides a map from node id to the revived
// objects.
//
// For objects that have to be revived but are not part of the object graph,
// this container provides a place where the objects can be stored so they are
// available to the runtime.
template <typename T>
class RevivedObjectContainer {
public:
// Insert an object that is not related to a node id. This usually means the
// object was not referenced by the object_graph, but is needed by other
// objects.
void Insert(std::unique_ptr<T> object) {
objects_.push_back(std::move(object));
}
// Insert an object that is tied to the given object graph node id.
void Insert(std::unique_ptr<T> object, int node_id) {
objects_by_id_[node_id] = object.get();
Insert(std::move(object));
}
// Find an object by the object graph node id.
// Returns nullptr if there is no such object.
T* Find(int node_id) {
auto it = objects_by_id_.find(node_id);
return it == objects_by_id_.end() ? nullptr : it->second;
}
private:
std::vector<std::unique_ptr<T>> objects_;
absl::flat_hash_map<int, T*> objects_by_id_;
};
// RevivedObjects is mainly used as a container for all the "state" owned by
// SavedModel. It stores all non-"user object" nodes from a SavedModel
// (https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/core/protobuf/saved_object_graph.proto#L57-L62)
// in a "fully constructed" state. It is effectively a strongly typed map, where
// each member is a map from the node id in the SavedObjectGraph's nodes
// (https://github.com/tensorflow/tensorflow/blob/568e2bef00f24af1159a0846abf67c099ca78a21/tensorflow/core/protobuf/saved_object_graph.proto#L25-L29)
// to the revived object of the corresponding type.
struct RevivedObjects {
// Order of declaration is important here: we want the RestoredResources to be
// freed after TFConcreteFunctions, for example.
gtl::FlatMap<int, std::unique_ptr<Variable>> variables;
gtl::FlatMap<int, std::unique_ptr<Asset>> assets;
gtl::FlatMap<int, std::unique_ptr<Constant>> constants;
gtl::FlatMap<int, std::unique_ptr<TFSignatureDefFunction>>
signature_def_functions;
RevivedObjectContainer<TFConcreteFunction> concrete_functions;
gtl::FlatMap<int, RestoredResource> restored_resources;
gtl::FlatMap<std::string, int> signatures_map;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_REVIVED_OBJECTS_H_
tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h
+49
View File
@@ -0,0 +1,49 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TENSORHANDLE_CONVERTIBLE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TENSORHANDLE_CONVERTIBLE_H_
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
namespace tensorflow {
// A common interface for objects that can be converted to a TensorHandle.
// Examples of objects that implement this include Variables, Constants, Assets,
// etc. This is used to convert captured objects into a ConcreteFunction's
// captured TensorHandles:
// https://github.com/tensorflow/tensorflow/blob/676a68963ea4b64fe479b9cede06aa8f5b290ab8/tensorflow/python/saved_model/load.py#L229-L240
class TensorHandleConvertible {
public:
explicit TensorHandleConvertible(ImmediateTensorHandlePtr handle)
: handle_(std::move(handle)) {}
ImmediateExecutionTensorHandle* handle() { return handle_.get(); }
// TensorHandleConvertible is movable, but not copyable.
TensorHandleConvertible(TensorHandleConvertible&& other) = default;
TensorHandleConvertible& operator=(TensorHandleConvertible&& other) = default;
virtual ~TensorHandleConvertible() = default;
protected:
TensorHandleConvertible(const TensorHandleConvertible&) = delete;
TensorHandleConvertible& operator=(const TensorHandleConvertible&) = delete;
ImmediateTensorHandlePtr handle_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TENSORHANDLE_CONVERTIBLE_H_
tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.cc
+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.
==============================================================================*/
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
TFConcreteFunction::TFConcreteFunction(std::unique_ptr<FlatTensorFunction> func,
FunctionMetadata metadata)
: func_(std::move(func)), metadata_(std::move(metadata)) {}
absl::Status TFConcreteFunction::Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
FunctionMetadata metadata, ImmediateExecutionContext* ctx,
std::unique_ptr<TFConcreteFunction>* out) {
std::unique_ptr<FlatTensorFunction> func;
TF_RETURN_IF_ERROR(FlatTensorFunction::Create(
function_def, std::move(captures), ctx, &func));
out->reset(new TFConcreteFunction(std::move(func), std::move(metadata)));
return absl::Status();
}
const FunctionMetadata& TFConcreteFunction::GetFunctionMetadata() const {
return metadata_;
}
absl::Status TFConcreteFunction::MakeCallOp(
absl::Span<AbstractTensorHandle* const> inputs, ImmediateOpPtr* out) const {
return func_->MakeCallOp(inputs, out);
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h
+85
View File
@@ -0,0 +1,85 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_H_
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// TF Eager Runtime-based implementation of a "ConcreteFunction" loaded from a
// saved model.
class TFConcreteFunction : public ConcreteFunction {
public:
// Factory function for creating a TFConcreteFunction.
//
// Params:
// function_def - The function_def associated with the created
// TFConcreteFunction. TFConcreteFunction will register this
// function_def with `ctx` on creation, and de-register it on
// destruction. function_def must be non-null, but
// otherwise has no lifetime requirements.
// captures - The captured TensorHandles associated with this
// TFConcreteFunction.
// metadata - The FunctionMetadata associated with this TFConcreteFunction.
// ctx - A handle to the Tensorflow runtime. This MUST be non-null and
// outlive TFConcreteFunction.
// out - The output TFConcreteFunction.
static absl::Status Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
FunctionMetadata metadata, ImmediateExecutionContext* ctx,
std::unique_ptr<TFConcreteFunction>* out);
// This method returns the "Call" Op used to execute the function.
absl::Status MakeCallOp(absl::Span<AbstractTensorHandle* const> inputs,
ImmediateOpPtr* out) const override;
const FunctionMetadata& GetFunctionMetadata() const override;
~TFConcreteFunction() override = default;
private:
TFConcreteFunction(std::unique_ptr<FlatTensorFunction> func,
FunctionMetadata metadata);
TFConcreteFunction(const TFConcreteFunction&) = delete;
TFConcreteFunction& operator=(const TFConcreteFunction&) = delete;
std::unique_ptr<FlatTensorFunction> func_;
FunctionMetadata metadata_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_H_
tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function_revival_state.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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_REVIVAL_STATE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_REVIVAL_STATE_H_
#include <memory>
#include <vector>
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// TFConcreteFunctionRevivalState wraps the state needed for building a
// TF_ConcreteFunction. This is mainly used in PartiallyRevivedObjects, which
// wraps partially constructed Function and Resource objects.
struct TFConcreteFunctionRevivalState {
// Index of the node in the SavedObjectGraph it was loaded from.
int node_id;
// Pointer to the original functiondef. fdef_ is guaranteed to be
// non-null.
const FunctionDef* fdef;
// TensorHandle captures for this funtion
std::vector<ImmediateExecutionTensorHandle*> captures;
// SavedConcreteFunction contains much of the metadata of the expected "types"
// of the inputs and outputs of a function.
// Note(bmzhao): saved_concrete_func_ is guaranteed to be non-null.
const SavedConcreteFunction* saved_concrete_func;
// This field is only present on TF2 ConcreteFunctions, and is useful for
// determining the original argument *names* of the function, (since the
// "canonicalized_input_signature" may append extra uniquifying integers).
// However, SavedBareConcreteFunctions do not have a FunctionSpec.
// Note(bmzhao): if function_spec_.has_value(), *function_spec_ is guaranteed
// to be non-null.
std::optional<const FunctionSpec*> function_spec;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_CONCRETE_FUNCTION_REVIVAL_STATE_H_
tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.cc
+65
View File
@@ -0,0 +1,65 @@
/* 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/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h"
#include <memory>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
TFSignatureDefFunction::TFSignatureDefFunction(
std::unique_ptr<FlatTensorFunction> func,
SignatureDefFunctionMetadata metadata)
: func_(std::move(func)), metadata_(std::move(metadata)) {}
absl::Status TFSignatureDefFunction::Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx,
std::unique_ptr<TFSignatureDefFunction>* out) {
std::unique_ptr<FlatTensorFunction> func;
TF_RETURN_IF_ERROR(FlatTensorFunction::Create(
function_def, std::move(captures), ctx, &func));
out->reset(new TFSignatureDefFunction(std::move(func), std::move(metadata)));
return absl::Status();
}
const SignatureDefFunctionMetadata&
TFSignatureDefFunction::GetFunctionMetadata() const {
return metadata_;
}
absl::Status TFSignatureDefFunction::MakeCallOp(
absl::Span<AbstractTensorHandle* const> inputs, ImmediateOpPtr* out) const {
return func_->MakeCallOp(inputs, out);
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/tf_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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_H_
#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// This is the TF eager runtime implementation of SignatureDefFunction (separate
// from the TFRT implementation). The user-facing API of SignatureDefFunctions
// and their semantic differences from ConcreteFunction are described here:
// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/cc/saved_model/experimental/public/signature_def_function.h#L30-L59
// Additional implementation notes are available here:
// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/c/experimental/saved_model/core/signature_def_function.h#L31-L48
class TFSignatureDefFunction : public SignatureDefFunction {
public:
// Factory function for creating a TFSignatureDefFunction.
//
// Params:
// function_def - The function_def associated with the created
// TFSignatureDefFunction. TFSignatureDefFunction will
// register this function_def with `ctx` on creation, and
// de-register it on destruction. function_def must be
// non-null, but otherwise has no lifetime requirements.
// captures - The captured TensorHandles associated with this
// TFConcreteFunction.
// metadata - FunctionMetadata associated with this TFSignatureDefFunction.
// ctx - A handle to the Tensorflow runtime. This MUST be non-null and
// outlive TFSignatureDefFunction.
// out - The output TFSignatureDefFunction.
static absl::Status Create(
const FunctionDef* function_def,
std::vector<ImmediateExecutionTensorHandle*> captures,
SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx,
std::unique_ptr<TFSignatureDefFunction>* out);
// This method creates a "Call" Op used to execute the function.
absl::Status MakeCallOp(absl::Span<AbstractTensorHandle* const> inputs,
ImmediateOpPtr* out) const override;
const SignatureDefFunctionMetadata& GetFunctionMetadata() const override;
~TFSignatureDefFunction() override = default;
private:
TFSignatureDefFunction(std::unique_ptr<FlatTensorFunction> func,
SignatureDefFunctionMetadata metadata);
TFSignatureDefFunction(const TFSignatureDefFunction&) = delete;
TFSignatureDefFunction& operator=(const TFSignatureDefFunction&) = delete;
std::unique_ptr<FlatTensorFunction> func_;
SignatureDefFunctionMetadata metadata_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_H_
tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function_revival_state.h
+55
View File
@@ -0,0 +1,55 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_REVIVAL_STATE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_REVIVAL_STATE_H_
#include <memory>
#include <vector>
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
// FunctionBuilder wraps the state needed for building a SignatureDefFunction.
// This is mainly used in PartiallyRevivedObjects, which wraps partially
// constructed Function and Resource objects.
struct TFSignatureDefFunctionRevivalState {
// Index of the node in the SavedObjectGraph it was loaded from.
int node_id = 0;
// Pointer to the original functiondef. fdef_ is guaranteed to be
// non-null.
const FunctionDef* fdef = nullptr;
// SavedConcreteFunction contains much of the metadata of the expected "types"
// of the inputs and outputs of a function.
// Note(bmzhao): saved_concrete_func_ is guaranteed to be non-null.
const SavedConcreteFunction* saved_concrete_func = nullptr;
// The name of the SignatureDef key.
std::string signature_key;
// TensorHandle captures for this funtion
std::vector<ImmediateExecutionTensorHandle*> captures;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_TF_SIGNATURE_DEF_FUNCTION_REVIVAL_STATE_H_
tensorflow/c/experimental/saved_model/core/revived_types/variable.cc
+130
View File
@@ -0,0 +1,130 @@
/* 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/c/experimental/saved_model/core/revived_types/variable.h"
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/ops/variable_ops.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
namespace tensorflow {
Variable::Variable(ImmediateExecutionContext* ctx, DataType dtype,
TensorShape shape, std::optional<std::string> name,
ImmediateTensorHandlePtr handle)
: TensorHandleConvertible(std::move(handle)),
name_(name.has_value() ? *name : "Variable"),
dtype_(dtype),
shape_(shape),
ctx_(ctx) {}
Variable::~Variable() {
// If the handle is null (perhaps because variable was std::moved from), then
// we don't have to do anything.
if (handle_ == nullptr) {
return;
}
absl::Status status = internal::DestroyResource(ctx_, handle_.get());
if (!status.ok()) {
LOG(ERROR) << "Error destroying variable: " << name_
<< "due to: " << status;
}
}
DataType Variable::dtype() { return dtype_; }
TensorShape Variable::shape() { return shape_; }
absl::Status Variable::Assign(ImmediateExecutionTensorHandle* handle) {
return internal::AssignVariable(ctx_, handle_.get(), dtype_, handle);
}
absl::Status Variable::ReadValue(ImmediateTensorHandlePtr* out) {
return internal::ReadVariable(ctx_, handle_.get(), dtype_, out);
}
absl::Status Variable::CreateUninitialized(
ImmediateExecutionContext* ctx, DataType dtype, TensorShape shape,
std::optional<std::string> name, const char* raw_device_name,
const std::vector<std::string>& component_devices,
std::unique_ptr<Variable>* output) {
ImmediateTensorHandlePtr handle;
if (component_devices.empty()) {
TF_RETURN_IF_ERROR(internal::CreateUninitializedResourceVariable(
ctx, dtype, shape, raw_device_name, &handle));
output->reset(
new Variable(ctx, dtype, shape, std::move(name), std::move(handle)));
return absl::Status();
}
if (!tensorflow::isa<EagerContext>(ctx)) {
return absl::InvalidArgumentError(
"Can only load distributed variables with EagerContext.");
}
EagerContext* eager_ctx = reinterpret_cast<EagerContext*>(ctx);
std::vector<TensorHandle*> handles;
for (const auto& device : component_devices) {
ImmediateTensorHandlePtr handlePtr;
TF_RETURN_IF_ERROR(internal::CreateUninitializedResourceVariable(
ctx, dtype, shape, device.empty() ? nullptr : device.c_str(),
&handlePtr));
if (!tensorflow::isa<TensorHandle>(handlePtr.get())) {
return absl::InternalError(
"Returned replica handle has unsupported type.");
}
handles.push_back(reinterpret_cast<TensorHandle*>(handlePtr.release()));
}
TensorHandle* packed_handle;
TF_RETURN_IF_ERROR(TensorHandle::CreatePackedHandle(
std::move(handles), eager_ctx, &packed_handle));
// The call to `CreatePackedHandle` incremented the handles' reference count,
// which we must now decrement to make the packed handle the owner of those
// handles. We can't loop through the `handles` vector because it was
// `std::move`d in the call above.
for (int i = 0; i != packed_handle->NumPackedHandles(); ++i) {
TensorHandle* component;
TF_RETURN_IF_ERROR(packed_handle->ExtractPackedHandle(i, &component));
component->Unref();
}
handle.reset(packed_handle);
output->reset(
new Variable(ctx, dtype, shape, std::move(name), std::move(handle)));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/revived_types/variable.h
+80
View File
@@ -0,0 +1,80 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_VARIABLE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_VARIABLE_H_
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
class Variable : public TensorHandleConvertible {
public:
// Creates an uninitialized resource variable. Note that a caller must
// call "assign" to associate a value with the variable.
static absl::Status CreateUninitialized(
ImmediateExecutionContext* ctx, DataType dtype, TensorShape shape,
std::optional<std::string> name, const char* raw_device_name,
const std::vector<std::string>& component_devices,
std::unique_ptr<Variable>* output);
// The dtype of the underlying variable.
DataType dtype();
// The shape of the underlying variable.
TensorShape shape();
// Updates the variable's contents with `handle`.
absl::Status Assign(ImmediateExecutionTensorHandle* handle);
// Reads the value of the variable, and stores it in `out`
absl::Status ReadValue(ImmediateTensorHandlePtr* out);
// Variable is movable, but not copyable.
Variable(Variable&& other) = default;
Variable& operator=(Variable&& other) = default;
~Variable() override;
private:
Variable(ImmediateExecutionContext* ctx, DataType dtype, TensorShape shape,
std::optional<std::string> name, ImmediateTensorHandlePtr handle);
Variable(const Variable& variable) = delete;
Variable& operator=(const Variable&) = delete;
std::string name_;
DataType dtype_;
TensorShape shape_;
// ctx_ must outlive Variable.
ImmediateExecutionContext* ctx_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_VARIABLE_H_
tensorflow/c/experimental/saved_model/core/saved_model_api.h
+66
View File
@@ -0,0 +1,66 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_
#include <memory>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
#include "tensorflow/cc/saved_model/bundle_v2.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Note(bmzhao): This class is only TEMPORARILY virtual, as a way to unblock
// TFRT integration with TF Serving. Do not add more virtual implementations of
// this class. Eventually we want to remove this virtual base class indirection
// and have only a single implementation.
class SavedModelAPI {
public:
// Retrieve a function from the TF2 SavedModel, using the "path" to a function
// in a TF2 savedmodel.
//
// Note: `function` is a double pointer, so that implementations are
// able to return a pointer to an internal member.
virtual absl::Status GetFunction(const std::string& function_path,
ConcreteFunction** function) = 0;
// Retrieve a list of child functions from a SavedModel given a starting node.
// 0 is the root node.
virtual absl::Status GetFunctions(
int node_id,
absl::flat_hash_map<std::string, ConcreteFunction*>* functions) = 0;
// Retrieve a SignatureDefFunction from a SavedModel, using the key of the
// SignatureDef map:
// https://github.com/tensorflow/tensorflow/blob/69b08900b1e991d84bce31f3b404f5ed768f339f/tensorflow/core/protobuf/meta_graph.proto#L89
virtual absl::Status GetSignatureDefFunction(
const std::string& signature_def_key,
SignatureDefFunction** function) = 0;
virtual SavedModelV2Bundle* GetBundle() = 0;
virtual ~SavedModelAPI() = default;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_API_H_
tensorflow/c/experimental/saved_model/core/saved_model_utils.cc
+541
View File
@@ -0,0 +1,541 @@
/* 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/c/experimental/saved_model/core/saved_model_utils.h"
#include <algorithm>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include "absl/strings/str_split.h"
#include "absl/types/optional.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/restored_resource_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_signature_def_function_revival_state.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/cc/saved_model/loader_util.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
#include "tensorflow/core/protobuf/trackable_object_graph.pb.h"
namespace tensorflow {
namespace internal {
namespace {
using StructuredValueDictEntry =
protobuf::MapPair<std::string, StructuredValue>;
// Maps from a Nodedef's name to its corresponding AttrValues, for a given
// Graphdef
using NodeAttrMap =
gtl::FlatMap<absl::string_view, const AttrValueMap*, StringPieceHasher>;
// Maps from a FunctionDef's name to FunctionDef, for a given FunctionDefLibrary
using FunctionDefMap =
gtl::FlatMap<absl::string_view, const tensorflow::FunctionDef*,
StringPieceHasher>;
// Looks up a SavedConstant's associated tensorproto from the NodeAttrMap and
// returns a tensorflow::Constant.
absl::Status ConstantFromSavedConstant(
ImmediateExecutionContext* ctx,
const tensorflow::SavedConstant& saved_constant,
const NodeAttrMap& node_attr_map, std::unique_ptr<Constant>* output) {
const std::string& const_op_name = saved_constant.operation();
const auto& node_name_and_attrs = node_attr_map.find(const_op_name);
if (node_name_and_attrs == node_attr_map.end()) {
return absl::FailedPreconditionError(
absl::StrCat("Unable to find Const operation with name'", const_op_name,
"' in SavedModel graphdef"));
}
const AttrValueMap* attrs = node_name_and_attrs->second;
const auto& attr_name_and_value = attrs->find("value");
if (attr_name_and_value == attrs->end()) {
return absl::FailedPreconditionError(
absl::StrCat("Unable to find Const operation '", const_op_name,
"'s value attribute"));
}
const TensorProto& tensor_proto = attr_name_and_value->second.tensor();
return internal::TensorProtoToConstant(ctx, tensor_proto, output);
}
// Perform some basic sanity checks on SavedConcreteFunction's input and
// output signatures with respect to the corresponding FunctionDef's input
// and output args.
absl::Status ValidateSavedFunctionCompatibleWithFunctionDef(
const SavedConcreteFunction& saved_concrete_function,
const FunctionDef* function_def) {
// tf.functions go through many transformations before becoming FunctionDefs
// 1. flatten user-provided inputs:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L2671-L2675
// 2. convert user-provided inputs to tensors:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L2687-L2688
// 3. filter any non-tensor, non-variable inputs:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L1840-L1841
// 4. concatenate any captured inputs:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L1912
// Since our API is limited to tf.functions annotated with input signatures,
// conditions 2 and 3 are trivially satisfied.
// We need to ensure that:
// flatten(input_signature).size() + captures.size() = fdef.signature().size()
// A concrete function's serialized "canonicalized_input_signature" comes
// from encoding its "structured_input_signature" field:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/saved_model/function_serialization.py#L70-L71
// The "structured_input_signature" is guaranteed to be a tuple of the python
// args, kwargs that correspond to the tf.function:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L1974-L1979
const std::string& name = function_def->signature().name();
const StructuredValue& input_signature =
saved_concrete_function.canonicalized_input_signature();
std::vector<const TensorSpecProto*> input_specs;
TF_RETURN_IF_ERROR(FlattenSignature(input_signature, &input_specs));
if (input_specs.size() + saved_concrete_function.bound_inputs_size() !=
function_def->signature().input_arg_size()) {
return absl::FailedPreconditionError(absl::StrCat(
"FunctionDef ", name, " has ",
function_def->signature().input_arg_size(),
" inputs, but the SavedConcreteFunction has ", input_specs.size(),
" flattened user inputs and ",
saved_concrete_function.bound_inputs_size(), " captured inputs."));
}
const StructuredValue& output_signature =
saved_concrete_function.output_signature();
std::vector<const TensorSpecProto*> output_specs;
TF_RETURN_IF_ERROR(FlattenSignature(output_signature, &output_specs));
if (output_specs.size() != function_def->signature().output_arg_size()) {
return absl::FailedPreconditionError(
absl::StrCat("FunctionDef ", name, " has ",
function_def->signature().output_arg_size(),
" outputs, but the SavedConcreteFunction has ",
output_specs.size(), " flattened outputs."));
}
return absl::Status();
}
} // namespace
absl::Status GetSignaturesMap(const SavedObjectGraph& saved_objects,
gtl::FlatMap<std::string, int>* signatures_map) {
if (saved_objects.nodes().empty()) {
return absl::FailedPreconditionError("Saved Object Graph was empty.");
}
const SavedObject& root = saved_objects.nodes(0);
const SavedObject* signatures = nullptr;
for (const auto& child : root.children()) {
if (child.local_name() == "signatures") {
if (child.node_id() >= saved_objects.nodes().size()) {
return absl::FailedPreconditionError(
absl::StrCat("Signature object had child node id ", child.node_id(),
" which exceeds the size of the set of nodes"));
}
signatures = &saved_objects.nodes(child.node_id());
}
}
// Some basic sanity checks that this object is actually our "signatures" map
if (signatures == nullptr) {
// This is where the "signatures" attribute is always set:
// https://github.com/tensorflow/tensorflow/blob/a2c542a0d83227568f9214a2af9a38ae3625976f/tensorflow/python/saved_model/save.py#L1106-L1109
return absl::FailedPreconditionError(
"SavedObjectGraph's root object must have a child 'signatures' object");
}
if (signatures->kind_case() != SavedObject::kUserObject) {
return absl::FailedPreconditionError(
"Signatures must be a SavedObject of type UserObject.");
}
if (signatures->user_object().identifier() != "signature_map") {
// This is where the string comes from:
// https://github.com/tensorflow/tensorflow/blob/c59af2913aaec235d883f50428efef1086f4c0e6/tensorflow/python/saved_model/signature_serialization.py#L220
return absl::FailedPreconditionError(
"Signatures SavedObject must have identifier 'signature_map'.");
}
for (const auto& child : signatures->children()) {
(*signatures_map)[child.local_name()] = child.node_id();
}
return absl::Status();
}
absl::Status ValidateSingleConcreteFunction(
const SavedFunction& saved_function) {
// We only allow loading functions that have an annotated input signature,
// which means there is 1:1 correspondence between tf.function
// <=> SavedFunction <=> SavedConcreteFunction <=> FunctionDef. This is
// the same restriction that MLIR has:
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc#L2677-L2707
if (saved_function.concrete_functions_size() != 1) {
return absl::FailedPreconditionError(
"Only tf.functions annotated with an input signature are supported "
"by SavedModelAPI. This means that there should only be a single "
"ConcreteFunction per tf.function");
}
return absl::Status();
}
absl::Status LoadSavedAsset(ImmediateExecutionContext* ctx,
const SavedAsset& asset,
const std::string& saved_model_dir,
absl::Span<const AssetFileDef> assets,
std::unique_ptr<Asset>* output) {
int asset_index = asset.asset_file_def_index();
if (asset_index >= assets.size()) {
return absl::FailedPreconditionError(absl::StrCat(
"SavedAsset contained asset index ", asset_index,
" but AssetFileDef only contains ", assets.size(), " # of assets"));
}
const std::string& asset_filename = assets[asset_index].filename();
return Asset::Create(ctx, saved_model_dir, asset_filename, output);
}
absl::Status TensorProtoToConstant(ImmediateExecutionContext* ctx,
const TensorProto& proto,
std::unique_ptr<Constant>* output) {
tensorflow::Tensor tensor;
bool parse_result = tensor.FromProto(proto);
if (!parse_result) {
return absl::InternalError("Failed to parse tensor from tensorproto");
}
TensorInterface tensor_interface(std::move(tensor));
return Constant::Create(ctx, &tensor_interface, output);
}
// This follows the python variable restoration logic:
// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/saved_model/load.py#L407
absl::Status LoadSavedVariable(ImmediateExecutionContext* ctx,
const SavedVariable& variable,
std::unique_ptr<Variable>* output) {
const std::string& name = variable.name();
tensorflow::TensorShape shape(variable.shape());
tensorflow::DataType dtype = variable.dtype();
std::vector<std::string> component_devices;
for (const auto& component :
variable.experimental_distributed_variable_components()) {
component_devices.push_back(component.device());
}
TF_RETURN_IF_ERROR(Variable::CreateUninitialized(
ctx, dtype, shape, name,
variable.device().empty() ? nullptr : variable.device().c_str(),
component_devices, output));
return absl::Status();
}
absl::Status LoadTFConcreteFunction(
const SavedConcreteFunction& saved_concrete_function,
const FunctionDef* function_def,
const std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>>&
captured_objects,
ImmediateExecutionContext* ctx, std::unique_ptr<TFConcreteFunction>* out) {
TF_RETURN_IF_ERROR(ValidateSavedFunctionCompatibleWithFunctionDef(
saved_concrete_function, function_def));
// Copy over captures
std::vector<ImmediateExecutionTensorHandle*> captures;
captures.reserve(saved_concrete_function.bound_inputs_size());
for (int bound_input : saved_concrete_function.bound_inputs()) {
auto iter = captured_objects.find(bound_input);
if (iter == captured_objects.end()) {
return absl::FailedPreconditionError(
absl::StrCat("Failed to find bound_input ", bound_input,
" for SavedConcreteFunction"));
}
captures.push_back(iter->second->handle());
}
return TFConcreteFunction::Create(function_def, std::move(captures), {}, ctx,
out);
}
absl::Status FlattenSignature(
const StructuredValue& signature,
std::vector<const TensorSpecProto*>* flattened_specs) {
// This follows the logic from
// https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/compiler/mlir/tensorflow/translate/import_model.cc#L2775
switch (signature.kind_case()) {
case StructuredValue::kDictValue: {
// Dictionaries must be sorted in order of keys
const DictValue& dict = signature.dict_value();
std::vector<const StructuredValueDictEntry*> entries;
entries.reserve(dict.fields_size());
for (const auto& field : dict.fields()) {
entries.push_back(&field);
}
std::sort(entries.begin(), entries.end(),
[](const StructuredValueDictEntry* x,
const StructuredValueDictEntry* y) {
return x->first < y->first;
});
for (const auto& entry : entries) {
TF_RETURN_IF_ERROR(FlattenSignature(entry->second, flattened_specs));
}
return absl::Status();
}
case StructuredValue::kTupleValue: {
const TupleValue& tuple = signature.tuple_value();
for (const StructuredValue& value : tuple.values()) {
TF_RETURN_IF_ERROR(FlattenSignature(value, flattened_specs));
}
return absl::Status();
}
case StructuredValue::kListValue: {
const ListValue& list = signature.list_value();
for (const StructuredValue& value : list.values()) {
TF_RETURN_IF_ERROR(FlattenSignature(value, flattened_specs));
}
return absl::Status();
}
case StructuredValue::kTensorSpecValue: {
flattened_specs->push_back(&signature.tensor_spec_value());
return absl::Status();
}
case StructuredValue::kNoneValue: {
// Base case: do nothing.
// This arises, for example, as the top-level object of an output
// signature when there are no return values.
return absl::Status();
}
default: {
return absl::InternalError(absl::StrCat(
"Unhandled structured value kind ", signature.kind_case()));
}
}
}
std::optional<int> FindNodeAtPath(absl::string_view path,
const SavedObjectGraph& object_graph) {
const auto& nodes = object_graph.nodes();
if (nodes.empty()) {
return std::nullopt;
}
// Starting from the root, iterate through the saved object graph, matching
// object names as we go.
int node_id = 0;
const SavedObject* current_node = &nodes.Get(node_id);
for (absl::string_view object_name : absl::StrSplit(path, '.')) {
auto child_node_iter = std::find_if(
current_node->children().begin(), current_node->children().end(),
[object_name](
const TrackableObjectGraph::TrackableObject::ObjectReference& obj) {
return object_name == obj.local_name();
});
if (child_node_iter == current_node->children().end()) {
return std::nullopt;
}
node_id = child_node_iter->node_id();
current_node = &nodes.Get(node_id);
}
return node_id;
}
gtl::FlatMap<absl::string_view, const AttrValueMap*, StringPieceHasher>
NodeToAttrMap(const tensorflow::GraphDef& graphdef) {
gtl::FlatMap<absl::string_view, const AttrValueMap*, StringPieceHasher>
result;
for (const tensorflow::NodeDef& node : graphdef.node()) {
result[node.name()] = &node.attr();
}
return result;
}
gtl::FlatMap<absl::string_view, const tensorflow::FunctionDef*,
StringPieceHasher>
FunctionNameToFunctionDefMap(const FunctionDefLibrary& library) {
gtl::FlatMap<absl::string_view, const tensorflow::FunctionDef*,
StringPieceHasher>
result;
for (const FunctionDef& function_def : library.function()) {
result[function_def.signature().name()] = &function_def;
}
return result;
}
absl::Status PartiallyReviveSavedModelObjects(
const MetaGraphDef& metagraph, ImmediateExecutionContext* context,
const std::string& directory, PartiallyRevivedObjects* objects) {
// This is needed to restore "Constant" nodes by looking up their
// "Value" attribute.
NodeAttrMap node_attr_map = NodeToAttrMap(metagraph.graph_def());
// These are needed for creating "Assets", by looking up their filenames.
std::vector<AssetFileDef> assets;
TF_RETURN_IF_ERROR(GetAssetFileDefs(metagraph, &assets));
// Signatures are needed for determining whether a function is a
// SignatureDefFunction or not.
gtl::FlatMap<std::string, int> signatures_map;
TF_RETURN_IF_ERROR(
GetSignaturesMap(metagraph.object_graph_def(), &signatures_map));
gtl::FlatMap<int, std::string> reversed_signatures_map;
reversed_signatures_map.reserve(signatures_map.size());
for (const auto& signature_key_and_node : signatures_map) {
reversed_signatures_map.emplace(signature_key_and_node.second,
signature_key_and_node.first);
}
// FunctionDefs are needed to help construct
// TFConcreteFunction/SignatureDefFunctions
const FunctionDefMap function_def_map =
internal::FunctionNameToFunctionDefMap(metagraph.graph_def().library());
// Iterate through all the saved objects, restoring objects (if we can) as we
// go. For objects that dependencies on other objects (resources/functions),
// we partially initialize "builders" that correspond to their currently known
// state, and gradually fill them out in subsequent passes.
for (int i = 0; i < metagraph.object_graph_def().nodes_size(); ++i) {
const SavedObject& node = metagraph.object_graph_def().nodes(i);
if (node.kind_case() == SavedObject::kVariable) {
std::unique_ptr<Variable> variable;
TF_RETURN_IF_ERROR(
LoadSavedVariable(context, node.variable(), &variable));
objects->variables[i] = std::move(variable);
} else if (node.kind_case() == SavedObject::kConstant) {
std::unique_ptr<Constant> constant;
TF_RETURN_IF_ERROR(ConstantFromSavedConstant(context, node.constant(),
node_attr_map, &constant));
objects->constants[i] = std::move(constant);
} else if (node.kind_case() == SavedObject::kAsset) {
std::unique_ptr<Asset> asset;
TF_RETURN_IF_ERROR(
LoadSavedAsset(context, node.asset(), directory, assets, &asset));
objects->assets[i] = std::move(asset);
} else if (node.kind_case() == SavedObject::kResource) {
RestoredResourceRevivalState resource_revival_state;
// We'll set the resource's functions in a subsequent pass, once we get
// all functions in a partially revived state.
resource_revival_state.device = node.resource().device();
objects->restored_resources[i] = std::move(resource_revival_state);
} else if (node.kind_case() == SavedObject::kFunction) {
// Get the SavedFunction node and skip if it has no concrete functions.
const SavedFunction& saved_function = node.function();
if (saved_function.concrete_functions_size() < 1) {
continue;
}
// Retrieve related function information.
const std::string& function_name = saved_function.concrete_functions(0);
const FunctionDef* function_def = function_def_map.at(function_name);
const SavedConcreteFunction& saved_concrete_func =
metagraph.object_graph_def().concrete_functions().at(function_name);
const FunctionSpec& function_spec = saved_function.function_spec();
// Construct either a SignatureDefFunctionBuilder or a
// ConcreteFunctionBuilder, depending on whether this node was a child
// of the "signatures" attribute from root object.
auto reverse_signature_iter = reversed_signatures_map.find(i);
if (reverse_signature_iter != reversed_signatures_map.end()) {
TFSignatureDefFunctionRevivalState func_revival_state;
func_revival_state.node_id = i;
func_revival_state.fdef = function_def;
func_revival_state.saved_concrete_func = &saved_concrete_func;
func_revival_state.signature_key = reverse_signature_iter->second;
objects->signature_def_functions[i] = std::move(func_revival_state);
} else {
TFConcreteFunctionRevivalState func_revival_state;
func_revival_state.node_id = i;
func_revival_state.fdef = function_def;
func_revival_state.saved_concrete_func = &saved_concrete_func;
func_revival_state.function_spec = &function_spec;
objects->concrete_functions[i] = std::move(func_revival_state);
}
} else if (node.kind_case() == SavedObject::kBareConcreteFunction) {
const SavedBareConcreteFunction& bare_cf = node.bare_concrete_function();
// Retrieve related function information.
const std::string& function_name = bare_cf.concrete_function_name();
const FunctionDef* function_def = function_def_map.at(function_name);
const SavedConcreteFunction& saved_concrete_func =
metagraph.object_graph_def().concrete_functions().at(function_name);
// Check whether this is a SignatureDefFunction, or not.
auto reverse_signature_iter = reversed_signatures_map.find(i);
if (reverse_signature_iter != reversed_signatures_map.end()) {
TFSignatureDefFunctionRevivalState func_revival_state;
func_revival_state.node_id = i;
func_revival_state.fdef = function_def;
func_revival_state.saved_concrete_func = &saved_concrete_func;
func_revival_state.signature_key = reverse_signature_iter->second;
objects->signature_def_functions[i] = std::move(func_revival_state);
} else {
TFConcreteFunctionRevivalState func_revival_state;
func_revival_state.node_id = i;
func_revival_state.fdef = function_def;
func_revival_state.saved_concrete_func = &saved_concrete_func;
objects->concrete_functions[i] = std::move(func_revival_state);
}
}
}
// Now that we've partially restored all functions, we can have resources
// point to them
for (auto& node_and_resource_revival_state : objects->restored_resources) {
int node_id = node_and_resource_revival_state.first;
const SavedObjectGraph& obj_graph = metagraph.object_graph_def();
const SavedObject& node = obj_graph.nodes(node_id);
RestoredResourceRevivalState& resource =
node_and_resource_revival_state.second;
for (const TrackableObjectGraph::TrackableObject::ObjectReference& child :
node.children()) {
int child_node_id = child.node_id();
// Note(bmzhao): The expected functions saved by a resource object are:
// "_create_resource", "_initialize", and "_destroy_resource".
// https://github.com/tensorflow/tensorflow/blob/ad66f588c1666ade8051feb42811fa27b285271c/tensorflow/python/training/tracking/tracking.py#L277-L281
if (child.local_name() == "_create_resource" &&
obj_graph.nodes(child.node_id()).kind_case() ==
SavedObject::kFunction) {
resource.create_resource = &objects->concrete_functions[child_node_id];
} else if (child.local_name() == "_initialize" &&
obj_graph.nodes(child.node_id()).kind_case() ==
SavedObject::kFunction) {
resource.initialize = &objects->concrete_functions[child_node_id];
} else if (child.local_name() == "_destroy_resource" &&
obj_graph.nodes(child.node_id()).kind_case() ==
SavedObject::kFunction) {
resource.destroy_resource = &objects->concrete_functions[child_node_id];
}
}
}
objects->signatures_map = std::move(signatures_map);
return absl::Status();
}
} // namespace internal
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/saved_model_utils.h
+120
View File
@@ -0,0 +1,120 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_UTILS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_UTILS_H_
// Some internal utility functions for the SavedModelAPI, factored out into a
// separately unit-testable header.
#include <memory>
#include <unordered_map>
#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/asset.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/partially_revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
namespace internal {
// Load a TensorProto into a tensorflow::Constant. This is similar to the
// constant loading logic in python:
// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/saved_model/load.py#L437
absl::Status TensorProtoToConstant(ImmediateExecutionContext* ctx,
const TensorProto& proto,
std::unique_ptr<Constant>* output);
// Creates a tensorflow::Variable from a SavedVariable. This is similar to the
// logic in:
// https://github.com/tensorflow/tensorflow/blob/516608035f85cec8b126712b0ff8407220206b22/tensorflow/python/saved_model/load.py#L407
// Note that the caller **must assign a value** to the loaded variable.
absl::Status LoadSavedVariable(ImmediateExecutionContext* ctx,
const SavedVariable& variable,
std::unique_ptr<Variable>* output);
absl::Status LoadSavedAsset(ImmediateExecutionContext* ctx,
const SavedAsset& asset,
const std::string& saved_model_dir,
absl::Span<const AssetFileDef> assets,
std::unique_ptr<Asset>* output);
// Creates a TFConcreteFunction from a SavedConcreteFunction.
absl::Status LoadTFConcreteFunction(
const SavedConcreteFunction& saved_concrete_function,
const FunctionDef* function_def,
const std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>>&
captured_objects,
ImmediateExecutionContext* ctx, std::unique_ptr<TFConcreteFunction>* out);
// Flattens `signature` into a vector of TensorSpecProto pointers back into
// `signature`. `signature` must outlive flattened_specs. `signature` must also
// be the input or output signature of a SavedConcreteFunction (i.e. "nested
// structures of tensorspecs").
absl::Status FlattenSignature(
const StructuredValue& signature,
std::vector<const TensorSpecProto*>* flattened_specs);
// Find the node id in `object_graph` at location `path`. `path` must be
// a dot-delimited string of object names relative to the root object. If no
// object is found, returns absl::nullopt.
std::optional<int> FindNodeAtPath(absl::string_view path,
const SavedObjectGraph& object_graph);
// Maps each node in `graphdef` to its corresponding Attribute Map.
// Callers must ensure that `graphdef` outlives the returned map.
gtl::FlatMap<absl::string_view, const AttrValueMap*, StringPieceHasher>
NodeToAttrMap(const tensorflow::GraphDef& graphdef);
// Maps the name of each FunctionDef in `library` to its corresponding
// FunctionDef. Callers must ensure `library` outlives the returned map.
gtl::FlatMap<absl::string_view, const tensorflow::FunctionDef*,
StringPieceHasher>
FunctionNameToFunctionDefMap(const FunctionDefLibrary& library);
// Finds the "signatures" object in the object graph, and fills a mapping of
// each signature's name to the corresponding function's node in the object
// graph.
absl::Status GetSignaturesMap(const SavedObjectGraph& saved_objects,
gtl::FlatMap<std::string, int>* signatures_map);
// Validates the `saved_function`.
absl::Status ValidateSingleConcreteFunction(
const SavedFunction& saved_function);
// Walks through the SavedObjectGraph in metagraph, and restores all nodes
// (except "UserDefinedObjects") with their corresponding type in
// "PartiallyRevivedObjects".
absl::Status PartiallyReviveSavedModelObjects(
const MetaGraphDef& metagraph, ImmediateExecutionContext* context,
const std::string& directory, PartiallyRevivedObjects* objects);
} // namespace internal
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SAVED_MODEL_UTILS_H_
tensorflow/c/experimental/saved_model/core/saved_variable_loading_test.cc
+170
View File
@@ -0,0 +1,170 @@
/* 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 <cstdint>
#include <memory>
#include <tuple>
#include <utility>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/common_runtime/eager/tensor_handle.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
namespace {
class SavedVariableLoadingTest
: public ::testing::TestWithParam<
std::tuple<DataType, std::vector<int64_t>>> {
public:
SavedVariableLoadingTest() {
SessionOptions options;
options.config.mutable_device_count()->insert({"CPU", 3});
std::vector<std::unique_ptr<Device>> devices;
TF_CHECK_OK(DeviceFactory::AddDevices(
options, "/job:localhost/replica:0/task:0", &devices));
device_mgr_ = absl::make_unique<StaticDeviceMgr>(std::move(devices));
ctx_ = testing::CreateTestingEagerContext(device_mgr_.get());
}
EagerContext* context() { return ctx_.get(); }
private:
std::unique_ptr<StaticDeviceMgr> device_mgr_;
EagerContextPtr ctx_;
};
// Sanity check that constructing a tensorflow::Variable from a SavedVariable
// 1. does not cause an error
// 2. preserves dtype and shape.
TEST_P(SavedVariableLoadingTest, LoadSavedVariableSuccessful) {
auto& test_params = GetParam();
DataType dtype = std::get<0>(test_params);
TensorShape shape(std::get<1>(test_params));
SavedVariable saved_variable;
saved_variable.set_dtype(dtype);
shape.AsProto(saved_variable.mutable_shape());
std::unique_ptr<Variable> var;
TF_EXPECT_OK(internal::LoadSavedVariable(context(), saved_variable, &var));
EXPECT_EQ(var->dtype(), dtype);
EXPECT_EQ(var->shape(), shape);
}
// Verify that a device specified in the SavedVariable is kept.
TEST_P(SavedVariableLoadingTest, LoadSavedVariableWithDevice) {
auto& test_params = GetParam();
DataType dtype = std::get<0>(test_params);
TensorShape shape(std::get<1>(test_params));
SavedVariable saved_variable;
saved_variable.set_dtype(dtype);
saved_variable.set_device("/job:localhost/replica:0/task:0/device:CPU:1"),
shape.AsProto(saved_variable.mutable_shape());
std::unique_ptr<Variable> var;
TF_ASSERT_OK(internal::LoadSavedVariable(context(), saved_variable, &var));
EXPECT_EQ(
absl::down_cast<TensorHandle*>(var->handle())->resource_device()->name(),
"/job:localhost/replica:0/task:0/device:CPU:1");
}
// Verify load failure if a non-existing device is specified.
TEST_P(SavedVariableLoadingTest, LoadSavedVariableWithInvalidDevice) {
auto& test_params = GetParam();
DataType dtype = std::get<0>(test_params);
TensorShape shape(std::get<1>(test_params));
SavedVariable saved_variable;
saved_variable.set_dtype(dtype);
saved_variable.set_device("/job:localhost/replica:0/task:0/device:CPU:99"),
shape.AsProto(saved_variable.mutable_shape());
std::unique_ptr<Variable> var;
ASSERT_NE(absl::OkStatus(),
internal::LoadSavedVariable(context(), saved_variable, &var));
}
// Assigning and reading values should yield
// consistent results.
TEST_P(SavedVariableLoadingTest, AssignAndReadVariableSuccesful) {
auto& test_params = GetParam();
DataType dtype = std::get<0>(test_params);
std::vector<int64_t> shape_vector = std::get<1>(test_params);
TensorShape shape(shape_vector);
// Create the variable.
absl::Status status;
std::unique_ptr<Variable> var;
TF_EXPECT_OK(Variable::CreateUninitialized(context(), dtype, shape,
std::nullopt, nullptr, {}, &var));
// Create a TensorHandle
ImmediateTensorHandlePtr expected_handle =
testing::CreateTensorHandle(context(), dtype, shape_vector, 42);
AbstractTensorPtr expected_tensor(expected_handle->Resolve(&status));
TF_EXPECT_OK(status) << status.message();
// Assign the tensorhandle to the variable.
TF_EXPECT_OK(var->Assign(expected_handle.get()));
// Read back the value from the variable
ImmediateTensorHandlePtr output_handle;
TF_EXPECT_OK(var->ReadValue(&output_handle));
AbstractTensorPtr output_tensor(output_handle->Resolve(&status));
TF_EXPECT_OK(status) << status.message();
// Check that output_tensor == expected_tensor
EXPECT_EQ(output_tensor->Type(), expected_tensor->Type());
EXPECT_EQ(output_tensor->NumElements(), expected_tensor->NumElements());
testing::CheckBufferDataIsEqual(
output_tensor->Type(), output_tensor->NumElements(),
output_tensor->Data(), expected_tensor->Data());
}
// Test against combinations of SavedVariables of
// 1. Varying dtypes
// 2. Varying shapes
INSTANTIATE_TEST_SUITE_P(
SavedVariableIntegerDtypesTest, SavedVariableLoadingTest,
::testing::Combine(
::testing::ValuesIn(testing::DataTypeSetToVector(kDataTypeIsInteger)),
::testing::ValuesIn(testing::InterestingShapes())));
INSTANTIATE_TEST_SUITE_P(
SavedVariableFloatingDtypesTest, SavedVariableLoadingTest,
::testing::Combine(::testing::Values(DT_FLOAT, DT_DOUBLE),
::testing::ValuesIn(testing::InterestingShapes())));
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/signature_def_function.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_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_H_
#include <memory>
#include <vector>
#include "absl/types/span.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
namespace tensorflow {
// See tensorflow/cc/experimental/saved_model/public/signature_def_function.h
// for SignatureDefFunction's intended user-facing semantics.
// This class is the "implementation" C++ part of the C++/C/C++ sandwich for
// a SignatureDefFunction.
// Note(bmzhao): Implementation-wise, SignatureDefFunctions are always saved as
// a "BareConcreteFunction", w/o a FunctionSpec, rather than a SavedFunction:
// https://github.com/tensorflow/tensorflow/blob/9bcefa44cd335c1db4a703a13da09f29ae1bbdb2/tensorflow/core/protobuf/saved_object_graph.proto#L60
// Additionally they are guaranteed to be children of the .signatures attribute
// of the root object, where the child object "name" is the signature_def key:
// https://github.com/tensorflow/tensorflow/blob/9bcefa44cd335c1db4a703a13da09f29ae1bbdb2/tensorflow/python/saved_model/signature_serialization.py#L181-L230
// One of the critical requirements of SignatureDef functions is that their
// inputs and outputs are "named". For example, a `.signatures` function:
// a. Requires users to pass: kwargs of all inputs:
// https://github.com/tensorflow/tensorflow/blob/26c4ee0c833e74f94d0102d8b005c41a28b44445/tensorflow/python/saved_model/signature_serialization.py#L119-L126
// b. Returns a dictionary of named outputs.
// https://github.com/tensorflow/tensorflow/blob/26c4ee0c833e74f94d0102d8b005c41a28b44445/tensorflow/python/saved_model/signature_serialization.py#L153-L161
// Since SignatureDefFunctions do not have FunctionSpecs, but guarantee the
// dictionary of inputs/outputs, we can parse these dictionaries' keys to obtain
// the input/output names of the SignatureDef:
// https://github.com/tensorflow/tensorflow/blob/9bcefa44cd335c1db4a703a13da09f29ae1bbdb2/tensorflow/core/protobuf/meta_graph.proto#L318-L321
class SignatureDefFunction {
public:
virtual ~SignatureDefFunction() = default;
// Creates a "Call" Op used to execute the function.
virtual absl::Status MakeCallOp(
absl::Span<AbstractTensorHandle* const> inputs,
ImmediateOpPtr* out) const = 0;
virtual const SignatureDefFunctionMetadata& GetFunctionMetadata() const = 0;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_H_
tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.cc
+42
View File
@@ -0,0 +1,42 @@
/* 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/c/experimental/saved_model/core/signature_def_function_metadata.h"
namespace tensorflow {
SignatureDefParam::SignatureDefParam(std::string name, TensorSpec spec)
: name_(std::move(name)), spec_(std::move(spec)) {}
const std::string& SignatureDefParam::name() const { return name_; }
const TensorSpec& SignatureDefParam::spec() const { return spec_; }
SignatureDefFunctionMetadata::SignatureDefFunctionMetadata(
std::vector<SignatureDefParam> arguments,
std::vector<SignatureDefParam> returns)
: arguments_(std::move(arguments)), returns_(std::move(returns)) {}
const std::vector<SignatureDefParam>& SignatureDefFunctionMetadata::arguments()
const {
return arguments_;
}
const std::vector<SignatureDefParam>& SignatureDefFunctionMetadata::returns()
const {
return returns_;
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h
+59
View File
@@ -0,0 +1,59 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_METADATA_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_METADATA_H_
#include <string>
#include <vector>
#include "tensorflow/c/experimental/saved_model/core/tensor_spec.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
// SignatureDefParam represents a named Tensor input or output to a
// SignatureDefFunction.
class SignatureDefParam {
public:
SignatureDefParam(std::string name, TensorSpec spec);
const std::string& name() const;
const TensorSpec& spec() const;
private:
std::string name_;
TensorSpec spec_;
};
class SignatureDefFunctionMetadata {
public:
SignatureDefFunctionMetadata() = default;
SignatureDefFunctionMetadata(std::vector<SignatureDefParam> arguments,
std::vector<SignatureDefParam> returns);
const std::vector<SignatureDefParam>& arguments() const;
const std::vector<SignatureDefParam>& returns() const;
private:
std::vector<SignatureDefParam> arguments_;
std::vector<SignatureDefParam> returns_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_SIGNATURE_DEF_FUNCTION_METADATA_H_
tensorflow/c/experimental/saved_model/core/signature_flattening_test.cc
+134
View File
@@ -0,0 +1,134 @@
/* 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 <string>
#include <vector>
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/tf_concrete_function_test_protos.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
namespace {
// Validates names, shapes, and dtypes of two tensorspecprotos are equivalent.
bool TensorSpecsAreEqual(const TensorSpecProto& spec,
const std::string& expected_name,
const PartialTensorShape& expected_shape,
DataType expected_dtype) {
return spec.name() == expected_name &&
PartialTensorShape(spec.shape()).IsIdenticalTo(expected_shape) &&
spec.dtype() == expected_dtype;
}
// This tests the common case for a tf.function w/o inputs. This ends up
// being serialized as a tuple of an empty tuple + empty dictionary
// (corresponding to the args, kwargs) of the function.
TEST(SignatureFlatteningTest, ZeroArgInputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::ZeroArgInputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 0);
}
// This tests the common case for a tf.function w/o outputs. This ends up
// being serialized as a "NoneValue".
TEST(SignatureFlatteningTest, ZeroRetOutputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::ZeroReturnOutputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 0);
}
TEST(SignatureFlatteningTest, SingleArgInputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::SingleArgInputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 1);
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[0],
/* expected_name = */ "x",
/* expected_shape = */ {1, 10},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[0]->DebugString();
}
TEST(SignatureFlatteningTest, SingleReturnOutputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::SingleReturnOutputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 1);
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[0],
/* expected_name = */ "",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[0]->DebugString();
}
TEST(SignatureFlatteningTest, ThreeArgInputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::ThreeArgInputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 3);
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[0],
/* expected_name = */ "x",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[0]->DebugString();
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[1],
/* expected_name = */ "y",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[1]->DebugString();
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[2],
/* expected_name = */ "z",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[2]->DebugString();
}
// This test has an exotic outputsignature of tuple of a
// dictionary<string,tensor>, tensor
TEST(SignatureFlatteningTest, ThreeReturnOutputSignature) {
std::vector<const TensorSpecProto*> flattened;
StructuredValue value = testing::ThreeReturnOutputSignature();
TF_EXPECT_OK(internal::FlattenSignature(value, &flattened));
EXPECT_EQ(flattened.size(), 3);
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[0],
/* expected_name = */ "0/a",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[0]->DebugString();
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[1],
/* expected_name = */ "0/b",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[1]->DebugString();
EXPECT_TRUE(TensorSpecsAreEqual(*flattened[2],
/* expected_name = */ "1",
/* expected_shape = */ {1},
/* expected_dtype = */ DT_FLOAT))
<< "Expected " << flattened[2]->DebugString();
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/tensor_spec.cc
+41
View File
@@ -0,0 +1,41 @@
/* 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/c/experimental/saved_model/core/tensor_spec.h"
#include <cstdint>
#include <initializer_list>
#include <utility>
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
TensorSpec::TensorSpec()
: shape_(std::initializer_list<int64_t>()), dtype_(DT_FLOAT) {}
TensorSpec::TensorSpec(PartialTensorShape shape, DataType dtype)
: shape_(std::move(shape)), dtype_(dtype) {}
TensorSpec::TensorSpec(const TensorSpecProto& proto)
: shape_(proto.shape()), dtype_(proto.dtype()) {}
const PartialTensorShape& TensorSpec::shape() const { return shape_; }
DataType TensorSpec::dtype() const { return dtype_; }
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/tensor_spec.h
+51
View File
@@ -0,0 +1,51 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_TENSOR_SPEC_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TENSOR_SPEC_H_
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
// Note(bmzhao): TensorSpec deliberately does not store the "name" from a
// TensorSpecProto. From edloper@, "Names should really be associated with
// parameters, not the tensors inside those parameters. This would be
// inconsistent with the corresponding Python class, but I don't think that's
// necessarily a problem. If it turns out later that we really need a name
// attribute here, we can always add it back in; but let's see how far we can
// get without it."
class TensorSpec {
public:
// Constructs a scalar, DT_FLOAT TensorSpec
TensorSpec();
TensorSpec(PartialTensorShape shape, DataType dtype);
explicit TensorSpec(const TensorSpecProto& proto);
const PartialTensorShape& shape() const;
DataType dtype() const;
private:
PartialTensorShape shape_;
DataType dtype_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TENSOR_SPEC_H_
tensorflow/c/experimental/saved_model/core/test_utils.cc
+157
View File
@@ -0,0 +1,157 @@
/* 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/c/experimental/saved_model/core/test_utils.h"
#include <memory>
#include <vector>
#include "absl/types/span.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/numeric_types.h"
#include "tensorflow/core/framework/register_types.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/bfloat16.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace testing {
std::unique_ptr<StaticDeviceMgr> CreateTestingDeviceMgr() {
return std::make_unique<StaticDeviceMgr>(
DeviceFactory::NewDevice("CPU", {}, "/job:localhost/replica:0/task:0"));
}
EagerContextPtr CreateTestingEagerContext(DeviceMgr* device_mgr) {
return EagerContextPtr(new EagerContext(
SessionOptions(),
tensorflow::ContextDevicePlacementPolicy::DEVICE_PLACEMENT_SILENT,
/* async= */ false, device_mgr,
/* device_mgr_owned= */ false, /* rendezvous= */ nullptr,
/* cluster_flr= */ nullptr));
}
std::vector<DataType> DataTypeSetToVector(DataTypeSet set) {
std::vector<DataType> result;
result.reserve(set.size());
for (DataType dt : set) {
result.push_back(dt);
}
return result;
}
std::vector<std::vector<int64_t>> InterestingShapes() {
std::vector<std::vector<int64_t>> interesting_shapes;
interesting_shapes.push_back({}); // Scalar
interesting_shapes.push_back({10}); // 1D Vector
interesting_shapes.push_back({3, 3}); // 2D Matrix
interesting_shapes.push_back({1, 4, 6, 10}); // Higher Dimension Tensor
return interesting_shapes;
}
ImmediateTensorHandlePtr CreateTensorHandle(ImmediateExecutionContext* ctx,
DataType dtype,
absl::Span<const int64_t> shape,
int8_t value) {
AbstractTensorPtr tensor(ctx->CreateTensor(dtype, shape));
CHECK_NE(tensor.get(), nullptr)
<< "Tensor creation failed for tensor of dtype: "
<< DataTypeString(dtype);
CHECK_EQ(tensor->Type(), dtype);
for (int i = 0; i < shape.size(); ++i) {
CHECK_EQ(tensor->Dim(i), shape[i]);
}
FillNumericTensorBuffer(tensor->Type(), tensor->NumElements(), tensor->Data(),
value);
ImmediateTensorHandlePtr handle(ctx->CreateLocalHandle(tensor.get()));
CHECK_NE(handle.get(), nullptr);
return handle;
}
void FillNumericTensorBuffer(DataType dtype, size_t num_elements, void* buffer,
int8_t value) {
switch (dtype) {
#define CASE(type) \
case DataTypeToEnum<type>::value: { \
type* typed_buffer = static_cast<type*>(buffer); \
for (size_t i = 0; i < num_elements; ++i) { \
typed_buffer[i] = static_cast<type>(value); \
} \
break; \
}
TF_CALL_INTEGRAL_TYPES(CASE);
TF_CALL_double(CASE);
TF_CALL_float(CASE);
TF_CALL_int4(CASE);
TF_CALL_uint4(CASE);
TF_CALL_int2(CASE);
TF_CALL_uint2(CASE);
#undef CASE
default:
CHECK(false) << "Unsupported data type: " << DataTypeString(dtype);
break;
}
}
// Checks the underlying data is equal for the buffers for two numeric tensors.
// Note: The caller must ensure to check that the dtypes and sizes of the
// underlying buffers are the same before calling this.
void CheckBufferDataIsEqual(DataType dtype, int64_t num_elements, void* a,
void* b) {
switch (dtype) {
#define CASE(type) \
case DataTypeToEnum<type>::value: { \
type* typed_a = static_cast<type*>(a); \
type* typed_b = static_cast<type*>(b); \
for (int64_t i = 0; i < num_elements; ++i) { \
if (DataTypeIsFloating(dtype)) { \
EXPECT_FLOAT_EQ(static_cast<float>(typed_a[i]), \
static_cast<float>(typed_b[i])); \
} else { \
EXPECT_EQ(typed_a[i], typed_b[i]); \
} \
} \
break; \
}
TF_CALL_INTEGRAL_TYPES(CASE);
TF_CALL_double(CASE);
TF_CALL_float(CASE);
TF_CALL_int4(CASE);
TF_CALL_uint4(CASE);
TF_CALL_int2(CASE);
TF_CALL_uint2(CASE);
#undef CASE
default:
CHECK(false) << "Unsupported data type: " << DataTypeString(dtype);
}
}
AbstractTensorPtr TensorHandleToTensor(ImmediateExecutionTensorHandle* handle) {
absl::Status status;
AbstractTensorPtr tensor(handle->Resolve(&status));
CHECK(status.ok()) << status.message();
CHECK_NE(tensor.get(), nullptr);
return tensor;
}
} // namespace testing
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/test_utils.h
+79
View File
@@ -0,0 +1,79 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_TEST_UTILS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TEST_UTILS_H_
#include <memory>
#include <vector>
#include "absl/types/span.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/core/common_runtime/device_mgr.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace testing {
// Creates a DeviceMgr suitable for local tests.
std::unique_ptr<StaticDeviceMgr> CreateTestingDeviceMgr();
// Creates an EagerContext suitable for local tests. Does not take ownership
// of `device_mgr`.
EagerContextPtr CreateTestingEagerContext(DeviceMgr* device_mgr);
// Converts a tensorflow::DatatypeSet to std::vector<DataType>.
// This is useful for tests using GTest's ::testing::ValuesIn, since
// DataTypeSet doesn't fullfill all the constraints of an STL-like iterable.
std::vector<DataType> DataTypeSetToVector(DataTypeSet set);
// Returns a vector of shapes intended to be "interesting" test cases.
// Currently, this returns scalar, 1D vector, 2D matrix, and a 4D tensor shapes
std::vector<std::vector<int64_t>> InterestingShapes();
// Returns a TensorHandle of `dtype` and `shape`, filled with `value`.
// `dtype` must be an integer dtype, float, or double.
// If a TensorHandle cannot be created successfully, this function will
// CHECK fail. This should only be used for testing purposes.
ImmediateTensorHandlePtr CreateTensorHandle(ImmediateExecutionContext* ctx,
DataType dtype,
absl::Span<const int64_t> shape,
int8_t value);
// Fills a numeric tensor's buffer with `value`.
// dtype must be any integer dtype, float or double.
void FillNumericTensorBuffer(DataType dtype, size_t num_elements, void* buffer,
int8_t value);
// Checks the underlying data is equal for the buffers for two numeric tensors.
// Note: The caller must ensure to check that the dtypes and sizes of the
// underlying buffers are the same before calling this.
// dtype must be any integer dtype, float, or double.
void CheckBufferDataIsEqual(DataType dtype, int64_t num_elements, void* a,
void* b);
// Converts a TensorHandle to a Tensor, and dies if unsuccessful. This should
// only be used for testing purposes.
AbstractTensorPtr TensorHandleToTensor(ImmediateExecutionTensorHandle* handle);
} // namespace testing
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TEST_UTILS_H_
tensorflow/c/experimental/saved_model/core/tf_concrete_function_loading_test.cc
+268
View File
@@ -0,0 +1,268 @@
/* 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 <cstdint>
#include <memory>
#include <string>
#include <unordered_map>
#include "absl/status/status.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"
#include "tensorflow/c/experimental/saved_model/core/tf_concrete_function_test_protos.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/error_codes.pb.h"
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
namespace tensorflow {
namespace {
class SavedConcreteFunctionLoadingTest : public ::testing::Test {
public:
SavedConcreteFunctionLoadingTest()
: device_mgr_(testing::CreateTestingDeviceMgr()),
ctx_(testing::CreateTestingEagerContext(device_mgr_.get())) {}
EagerContext* context() { return ctx_.get(); }
private:
std::unique_ptr<StaticDeviceMgr> device_mgr_;
EagerContextPtr ctx_;
};
class DummyCapture : public TensorHandleConvertible {
public:
DummyCapture(ImmediateExecutionContext* ctx, int8_t value)
: TensorHandleConvertible(
testing::CreateTensorHandle(ctx, DT_FLOAT, {2, 4}, value)) {}
};
FunctionDef FuncDefWithNumInputsOutputs(int num_inputs, int num_outputs) {
FunctionDef func;
OpDef* signature = func.mutable_signature();
for (int i = 0; i < num_inputs; ++i) {
signature->add_input_arg();
}
for (int i = 0; i < num_outputs; ++i) {
signature->add_output_arg();
}
return func;
}
// A SavedConcreteFunction whose canonicalized input signature
// has less inputs than its corresponding FunctionDef should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest, TooFewInputsInSavedConcreteFunction) {
// `saved` has 1 input
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::SingleArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
// `func` has 2 inputs
FunctionDef func = FuncDefWithNumInputsOutputs(2, 0);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status =
internal::LoadTFConcreteFunction(saved, &func, {}, context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose canonicalized input signature length +
// captures is less than its corresponding FunctionDef should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest,
TooFewInputsWithCapturesInSavedConcreteFunction) {
// `saved` has 1 input, and 1 capture, for a total of 2 inputs
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::SingleArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
saved.add_bound_inputs(5);
// `func` has 3 inputs
FunctionDef func = FuncDefWithNumInputsOutputs(3, 0);
std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>> captures;
captures[5] = std::make_unique<DummyCapture>(context(), 10);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status = internal::LoadTFConcreteFunction(saved, &func, captures,
context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose canonicalized input signature
// has more inputs than its corresponding FunctionDef should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest, TooManyInputsInSavedConcreteFunction) {
// `saved` has 3 inputs
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::ThreeArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
// `func` has 2 inputs
FunctionDef func = FuncDefWithNumInputsOutputs(2, 0);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status =
internal::LoadTFConcreteFunction(saved, &func, {}, context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose canonicalized input signature
// has the same number of inputs than its corresponding FunctionDef, but has
// additional captures should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest,
TooManyInputsWithCaptureInSavedConcreteFunction) {
// `saved` has 3 inputs, and 1 capture, for a total of 4 inputs.
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::ThreeArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
saved.add_bound_inputs(5);
// `func` has 3 inputs.
FunctionDef func = FuncDefWithNumInputsOutputs(3, 0);
std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>> captures;
captures[5] = std::make_unique<DummyCapture>(context(), 10);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status = internal::LoadTFConcreteFunction(saved, &func, captures,
context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose capture refers to an index not in the capture
// map should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest, ImproperCaptureIndex) {
// `saved` has 3 inputs, 1 capture, for a total of 4 inputs
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::ThreeArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
// Capture is at index "10"
saved.add_bound_inputs(10);
// `func` has 4 inputs
FunctionDef func = FuncDefWithNumInputsOutputs(4, 0);
// `captures` only has a capture for index 5
std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>> captures;
captures[5] = std::make_unique<DummyCapture>(context(), 10);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status = internal::LoadTFConcreteFunction(saved, &func, captures,
context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose outputs are fewer than its corresponding
// functiondef should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest, TooFewOutputsInSavedConcreteFunction) {
// `saved` has 0 inputs, 1 output
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::ZeroArgInputSignature();
*saved.mutable_output_signature() = testing::SingleReturnOutputSignature();
// `func` has 0 inputs, 2 outputs
FunctionDef func = FuncDefWithNumInputsOutputs(0, 2);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status =
internal::LoadTFConcreteFunction(saved, &func, {}, context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose outputs exceed its corresponding functiondef
// should cause an error.
TEST_F(SavedConcreteFunctionLoadingTest,
TooManyOutputsInSavedConcreteFunction) {
// `saved` has 1 input, 3 outputs
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::SingleArgInputSignature();
*saved.mutable_output_signature() = testing::ThreeReturnOutputSignature();
// `func` has 1 input, 2 outputs
FunctionDef func = FuncDefWithNumInputsOutputs(1, 2);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status =
internal::LoadTFConcreteFunction(saved, &func, {}, context(), &result);
EXPECT_EQ(status.code(), error::FAILED_PRECONDITION) << status.message();
}
// A SavedConcreteFunction whose (inputs + captures) = functiondef inputs,
// and whose outputs = functiondef outputs should successfully load.
TEST_F(SavedConcreteFunctionLoadingTest, SuccessfulLoad) {
// `saved` has 1 input, 2 captures, 3 outputs
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::SingleArgInputSignature();
*saved.mutable_output_signature() = testing::ThreeReturnOutputSignature();
saved.add_bound_inputs(2);
saved.add_bound_inputs(5);
// `func` has 3 inputs, 3 outputs
FunctionDef func = FuncDefWithNumInputsOutputs(3, 3);
std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>> captures;
captures[2] = std::make_unique<DummyCapture>(context(), 1);
captures[5] = std::make_unique<DummyCapture>(context(), 10);
std::unique_ptr<TFConcreteFunction> result;
absl::Status status = internal::LoadTFConcreteFunction(saved, &func, captures,
context(), &result);
TF_EXPECT_OK(status) << status.message();
}
// A TFConcreteFunction should register functiondefs on creation, and
// remove them upon deletion.
TEST_F(SavedConcreteFunctionLoadingTest, RegistersAndRemovesFunctionDefs) {
std::string func_name = "FooBarBazWombatFunction";
SavedConcreteFunction saved;
*saved.mutable_canonicalized_input_signature() =
testing::ZeroArgInputSignature();
*saved.mutable_output_signature() = testing::ZeroReturnOutputSignature();
FunctionDef func = FuncDefWithNumInputsOutputs(0, 0);
*func.mutable_signature()->mutable_name() = func_name;
{
std::unique_ptr<TFConcreteFunction> result;
absl::Status status =
internal::LoadTFConcreteFunction(saved, &func, {}, context(), &result);
TF_EXPECT_OK(status) << status.message();
// The function should be registered with context.
EXPECT_TRUE(context()->FindFunctionByName(func_name));
}
// After `result's` destructor runs, the function should no longer be
// registered with context.
EXPECT_FALSE(context()->FindFunctionByName(func_name));
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/tf_concrete_function_test_protos.cc
+212
View File
@@ -0,0 +1,212 @@
/* 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/c/experimental/saved_model/core/tf_concrete_function_test_protos.h"
#include <string>
#include "absl/strings/string_view.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
namespace testing {
namespace {
constexpr absl::string_view kZeroArgInputSignatureTextProto = R"(
tuple_value: {
values: {
tuple_value: {
}
}
values: {
dict_value: {
}
}
}
)";
constexpr absl::string_view kSingleArgInputSignatureTextProto = R"(
tuple_value: {
values: {
tuple_value: {
values: {
tensor_spec_value: {
name : "x"
shape: {
dim: {
size: 1
}
dim: {
size: 10
}
}
dtype: DT_FLOAT
}
}
}
}
values: {
dict_value: {
}
}
}
)";
constexpr absl::string_view kThreeArgInputSignatureTextProto = R"(
tuple_value: {
values: {
tuple_value: {
values: {
tensor_spec_value: {
name : "x"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
values: {
tensor_spec_value: {
name : "y"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
values: {
tensor_spec_value: {
name : "z"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
}
}
values: {
dict_value: {
}
}
}
)";
constexpr absl::string_view kZeroReturnOutputSignatureTextProto = R"(
none_value: {}
)";
constexpr absl::string_view kSingleReturnOutputSignatureTextProto = R"(
tensor_spec_value: {
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
)";
constexpr absl::string_view kThreeReturnOutputSignatureTextProto = R"(
tuple_value: {
values: {
dict_value: {
fields: {
key : "a"
value: {
tensor_spec_value: {
name : "0/a"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
}
fields: {
key : "b"
value: {
tensor_spec_value: {
name : "0/b"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
}
}
}
values: {
tensor_spec_value: {
name : "1"
shape: {
dim: {
size: 1
}
}
dtype: DT_FLOAT
}
}
}
)";
StructuredValue ParseStructuredValue(absl::string_view text_proto) {
StructuredValue value;
CHECK(tensorflow::protobuf::TextFormat::ParseFromString(text_proto, &value));
return value;
}
} // namespace
StructuredValue ZeroArgInputSignature() {
return ParseStructuredValue(kZeroArgInputSignatureTextProto);
}
StructuredValue SingleArgInputSignature() {
return ParseStructuredValue(kSingleArgInputSignatureTextProto);
}
StructuredValue ThreeArgInputSignature() {
return ParseStructuredValue(kThreeArgInputSignatureTextProto);
}
StructuredValue ZeroReturnOutputSignature() {
return ParseStructuredValue(kZeroReturnOutputSignatureTextProto);
}
StructuredValue SingleReturnOutputSignature() {
return ParseStructuredValue(kSingleReturnOutputSignatureTextProto);
}
StructuredValue ThreeReturnOutputSignature() {
return ParseStructuredValue(kThreeReturnOutputSignatureTextProto);
}
} // namespace testing
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/tf_concrete_function_test_protos.h
+50
View File
@@ -0,0 +1,50 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_CONCRETE_FUNCTION_TEST_PROTOS_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_CONCRETE_FUNCTION_TEST_PROTOS_H_
#include "tensorflow/core/protobuf/struct.pb.h"
namespace tensorflow {
namespace testing {
// Returns a StructuredValue corresponding to the serialized InputSignature of a
// tf.function with 0 inputs
StructuredValue ZeroArgInputSignature();
// Returns a StructuredValue corresponding to the serialized InputSignature of a
// tf.function with 1 input
StructuredValue SingleArgInputSignature();
// Returns a StructuredValue corresponding to the serialized InputSignature of a
// tf.function with 3 inputs
StructuredValue ThreeArgInputSignature();
// Returns a StructuredValue corresponding to the serialized OutputSignature of
// a tf.function with no return values
StructuredValue ZeroReturnOutputSignature();
// Returns a StructuredValue corresponding to the serialized OutputSignature of
// a tf.function with a single tensor output
StructuredValue SingleReturnOutputSignature();
// Returns a StructuredValue corresponding to the serialized OutputSignature of
// a tf.function with three tensor outputs
StructuredValue ThreeReturnOutputSignature();
} // namespace testing
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_CONCRETE_FUNCTION_TEST_PROTOS_H_
tensorflow/c/experimental/saved_model/core/tf_saved_model_api.cc
+322
View File
@@ -0,0 +1,322 @@
/* 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/c/experimental/saved_model/core/tf_saved_model_api.h"
#include <memory>
#include <string>
#include <unordered_set>
#include <utility>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/ops/restore_ops.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/partially_revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_utils.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
#include "tensorflow/cc/saved_model/bundle_v2.h"
#include "tensorflow/cc/saved_model/constants.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_util.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/gtl/flatmap.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/macros.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/tstring.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/protobuf/trackable_object_graph.pb.h"
namespace tensorflow {
// Maps from a FunctionDef's name to FunctionDef, for a given FunctionDefLibrary
using FunctionDefMap =
gtl::FlatMap<absl::string_view, const tensorflow::FunctionDef*,
StringPieceHasher>;
// Maps from a functiondef's name to the corresponding "TFConcreteFunction"
using FlatTensorFunctionMap =
gtl::FlatMap<std::string, std::unique_ptr<FlatTensorFunction>>;
namespace {
const TrackableObjectGraph::TrackableObject::SerializedTensor*
FindSerializedTensorInTrackable(
const TrackableObjectGraph::TrackableObject& trackable_object,
absl::string_view name) {
for (const auto& maybe_serialized_tensor : trackable_object.attributes()) {
if (maybe_serialized_tensor.name() == name) {
return &maybe_serialized_tensor;
}
}
return nullptr;
}
// This function reads the Checkpoint embedded in the SavedModel, and calls the
// appropriate Restore ops on each of the variables.
// Note(bmzhao): Conceptually, objects that contain checkpointable state
// implement the "_gather_saveables_for_checkpoint" method
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/python/training/tracking/base.py#L953-L983
// which returns a dict of string key -> EITHER:
// 1. python callable (taking a checkpoint key) returning SaveableObject OR
// 2. variable (partitioned/resource/reference or otherwise)
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/python/training/saving/saveable_object.py#L58.
// The string key becomes the "name" attribute of the SerializedTensor proto
// in the TrackableObjectGraph,
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/core/protobuf/trackable_object_graph.proto#L26
// And the checkpoint_key is a globally unique string derived from this name:
// https://github.com/tensorflow/tensorflow/blob/842df9e6b516e42578a8d23b35d41176b9a6cf1d/tensorflow/python/training/tracking/graph_view.py#L236-L241
// SaveableObjects model the information needed to pass to the SaveV2/RestoreV2
// ops via their SaveSpec members
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/python/training/saving/saveable_object.py#L21,
// which contain the "real" checkpoint keys into the TensorBundle SSTable.
// They also contain the logic needed to take the restored tensors from
// RestoreV2 and load them back into the "object" they came from via their
// overridden "restore" method:
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/python/training/saving/saveable_object.py#L85
absl::Status RestoreCheckpoint(SavedModelV2Bundle* bundle,
const RevivedObjects& revived_objects,
const std::string& directory,
ImmediateExecutionContext* context) {
// TODO(bmzhao): Batch up all the restores into a single restore op per
// device, following logic in MultiDeviceSaver.
TF_RETURN_IF_ERROR(bundle->VisitObjectsToRestore(
[&revived_objects, &directory, context, bundle](
int node, const TrackableObjectGraph::TrackableObject& trackable) {
if (bundle->saved_object_graph().nodes(node).kind_case() !=
SavedObject::kVariable) {
// TODO(bmzhao): This requires using the newly added Save/Restore
// functions from
// https://github.com/tensorflow/tensorflow/commit/df6b21c13c82b5d0981642cfe18f10e60f78ea5c
LOG(WARNING) << "Restoring non-variable objects has not been "
"implemented yet. (Kind="
<< bundle->saved_object_graph().nodes(node).kind_case()
<< ")";
return absl::OkStatus();
}
Variable* variable = revived_objects.variables.at(node).get();
// Restore the tensor's value from the checkpoint
const TrackableObjectGraph::TrackableObject::SerializedTensor*
attribute =
FindSerializedTensorInTrackable(trackable, "VARIABLE_VALUE");
if (attribute == nullptr) {
return absl::FailedPreconditionError(
"Could not find SerializedTensor with name VARIABLE_VALUE for "
"saved variable");
}
const std::string& checkpoint_key = attribute->checkpoint_key();
if (!bundle->variable_reader()->Contains(checkpoint_key)) {
LOG(WARNING) << "No checkpoint entry found for " << checkpoint_key
<< ". Variable will be uninitialized.";
return absl::Status();
}
std::string variables_path_prefix =
io::JoinPath(directory, kSavedModelVariablesDirectory,
kSavedModelVariablesFilename);
ImmediateTensorHandlePtr restored_output;
TF_RETURN_IF_ERROR(internal::SingleRestore(
context, variables_path_prefix, checkpoint_key, variable->dtype(),
&restored_output));
// Assign the restored tensor's value to the variable
return variable->Assign(restored_output.get());
}));
return absl::Status();
}
absl::Status InitializeAllResources(const RevivedObjects& revived) {
for (const auto& node_and_resource : revived.restored_resources) {
const RestoredResource& resource = node_and_resource.second;
TF_RETURN_IF_ERROR(resource.Initialize());
}
return absl::Status();
}
} // namespace
absl::Status TFSavedModelAPI::GetFunction(const std::string& function_path,
ConcreteFunction** function) {
std::optional<int> node =
internal::FindNodeAtPath(function_path, bundle_.saved_object_graph());
if (!node.has_value()) {
return absl::NotFoundError(
absl::StrCat("No saved object found at path ", function_path));
}
*function = revived_objects_.concrete_functions.Find(*node);
if (*function == nullptr) {
return absl::NotFoundError(
absl::StrCat("No function found at path ", function_path));
}
return absl::Status();
}
absl::Status TFSavedModelAPI::GetFunctions(
int node_id,
absl::flat_hash_map<std::string, ConcreteFunction*>* functions) {
const auto& nodes = bundle_.saved_object_graph().nodes();
if (node_id >= nodes.size()) {
return absl::OutOfRangeError(
absl::StrCat("node_id ", node_id,
" not found. Maximum node ID: ", nodes.size() - 1));
}
const SavedObject* current_node = &nodes.Get(node_id);
for (const auto& child : current_node->children()) {
ConcreteFunction* concrete_fn;
absl::Status status = GetFunction(child.local_name(), &concrete_fn);
if (status.ok()) {
(*functions)[child.local_name()] = concrete_fn;
}
}
return absl::Status();
}
absl::Status TFSavedModelAPI::GetSignatureDefFunction(
const std::string& signature_def_key, SignatureDefFunction** function) {
auto signatures_iter =
revived_objects_.signatures_map.find(signature_def_key);
if (signatures_iter == revived_objects_.signatures_map.end()) {
return absl::NotFoundError(absl::StrCat("No signature with key ",
signature_def_key, " was found"));
}
int node = signatures_iter->second;
auto function_iter = revived_objects_.signature_def_functions.find(node);
if (function_iter == revived_objects_.signature_def_functions.end()) {
return absl::InternalError(
absl::StrCat("Unable to find SignatureDefFunction associated with key ",
signature_def_key, " despite key being valid."));
}
*function = function_iter->second.get();
return absl::Status();
}
absl::Status TFSavedModelAPI::GetVariable(const std::string& variable_path,
Variable** variable) {
std::optional<int> node =
internal::FindNodeAtPath(variable_path, bundle_.saved_object_graph());
if (!node.has_value()) {
return absl::NotFoundError(
absl::StrCat("No saved object found at path ", variable_path));
}
auto variables_iter = revived_objects_.variables.find(*node);
if (variables_iter == revived_objects_.variables.end()) {
return absl::NotFoundError(
absl::StrCat("No variable found at path ", variable_path));
}
*variable = variables_iter->second.get();
return absl::Status();
}
SavedModelV2Bundle* TFSavedModelAPI::GetBundle() { return &this->bundle_; }
TFSavedModelAPI::TFSavedModelAPI(const std::string& directory,
SavedModelV2Bundle bundle,
RevivedObjects revived_objects)
: directory_(directory),
bundle_(std::move(bundle)),
revived_objects_(std::move(revived_objects)) {}
absl::Status TFSavedModelAPI::Load(
const std::string& directory,
const std::optional<std::unordered_set<std::string>>& tags,
ImmediateExecutionContext* context, std::unique_ptr<TFSavedModelAPI>* out) {
// TODO(bmzhao): Add support for loading a TF1 SavedModel.
if (tags) {
return absl::UnimplementedError(
"Loading saved models with explicit tags will be supported in the "
"future");
}
SavedModelV2Bundle bundle;
TF_RETURN_IF_ERROR(SavedModelV2Bundle::Load(directory, &bundle));
// TODO(bmzhao): Mangle loaded function names so that different
// models loaded in the same runtime Context don't clobber eachother.
// This occurs in python here:
// https://github.com/tensorflow/tensorflow/blob/285b5fa15405c5e2c084080f52a1818be8648079/tensorflow/python/saved_model/function_deserialization.py#L438-L454
// For each node in the graph, we should initialize an object of the
// corresponding type. For objects that depend on the initialization of other
// objects (like functions which capture resources), we will initialize them
// later.
PartiallyRevivedObjects partially_revived_objects;
TF_RETURN_IF_ERROR(internal::PartiallyReviveSavedModelObjects(
bundle.meta_graph_def(), context, directory, &partially_revived_objects));
RevivedObjects revived_objects;
TF_RETURN_IF_ERROR(partially_revived_objects.Build(
context, bundle.saved_object_graph(), &revived_objects));
// Revive function library functions as concrete functions without captures.
// This is necessary because object graph functions may refer to functions
// _not_ in the object graph: A while loop, for example, will create two
// auxiliary `while_cond` and `while_body` functions that are only present in
// the graph def function library.
for (const FunctionDef& function :
bundle.meta_graph_def().graph_def().library().function()) {
std::unique_ptr<TFConcreteFunction> concrete_function;
TF_RETURN_IF_ERROR(TFConcreteFunction::Create(/*function_def=*/&function,
/*captures=*/{},
/*metadata=*/{},
/*ctx=*/context,
/*out=*/&concrete_function));
revived_objects.concrete_functions.Insert(std::move(concrete_function));
}
TF_RETURN_IF_ERROR(
RestoreCheckpoint(&bundle, revived_objects, directory, context));
TF_RETURN_IF_ERROR(InitializeAllResources(revived_objects));
out->reset(new TFSavedModelAPI(directory, std::move(bundle),
std::move(revived_objects)));
return absl::Status();
}
} // namespace tensorflow
tensorflow/c/experimental/saved_model/core/tf_saved_model_api.h
+94
View File
@@ -0,0 +1,94 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SAVED_MODEL_API_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SAVED_MODEL_API_H_
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/revived_objects.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_api.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
#include "tensorflow/cc/saved_model/bundle_v2.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// An implementation of the SavedModelAPI using the TF Eager runtime. See
// https://github.com/tensorflow/community/blob/master/rfcs/20200218-tf-c-saved-model.md
// Conceptually, there are many differences between a tf.function and
// a FunctionDef is executed by the C API.
// 1. A tf.function is polymorphic, meaning it can correspond to multiple
// ConcreteFunctions (of differing shapes, python arguments, etc). A
// FunctionDef corresponds to a single ConcreteFunction.
// 2. A tf.function can take arbitrary python inputs, whereas the FunctionDef
// only accepts tensors.
// 3. A tf.function is a closure that can contain captured inputs, whereas
// FunctionDefs loaded from SavedModels are "functional" (all inputs are
// explicitly passed as arguments).
// The SavedModelAPI only supports loading tf.functions annotated with input
// signatures so that we ensure that there is a 1:1 mapping between tf.function
// -> FunctionDef, and have a guarantee that all inputs are tensors.
// (https://github.com/tensorflow/tensorflow/blob/2b96f3662bd776e277f86997659e61046b56c315/tensorflow/python/eager/def_function.py#L1167-L1171),
class TFSavedModelAPI : public SavedModelAPI {
public:
absl::Status GetFunction(const std::string& function_path,
ConcreteFunction** function) override;
absl::Status GetFunctions(
int node_id,
absl::flat_hash_map<std::string, ConcreteFunction*>* functions) override;
absl::Status GetSignatureDefFunction(
const std::string& signature_def_key,
SignatureDefFunction** function) override;
static absl::Status Load(
const std::string& directory,
const std::optional<std::unordered_set<std::string>>& tags,
ImmediateExecutionContext* context,
std::unique_ptr<TFSavedModelAPI>* out);
~TFSavedModelAPI() override = default;
absl::Status GetVariable(const std::string& variable_path,
Variable** variable);
SavedModelV2Bundle* GetBundle() override;
private:
TFSavedModelAPI(const std::string& directory, SavedModelV2Bundle bundle,
RevivedObjects revived_objects);
std::string directory_;
SavedModelV2Bundle bundle_;
RevivedObjects revived_objects_;
};
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SAVED_MODEL_API_H_
tensorflow/c/experimental/saved_model/internal/BUILD
+382
View File
@@ -0,0 +1,382 @@
# Experimental Implementation of SavedModel C APIs for TensorFlow. See RFC
# https://github.com/tensorflow/community/pull/207
# External clients should not worry about this directory; all contents are implementation details.
# Code in this directory is intended to form the glue between the C API and the internal C++
# implementation by
# 1. mapping C API calls onto correponding methods of C++ objects
# 2. mapping opaque C types onto C++ classes
# Note(bmzhao): The *.cc files in this directory form the direct implementation of the
# C API functions exposed in tf/c/experimental/saved_model/public/.
# Note(bmzhao): All *type.h files in this directory are the internal definitions of
# the opaque C types. These headers should only be visible to internal tensorflow
# implementors.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
"tf_copts",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "concrete_function",
srcs = [
"concrete_function.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:concrete_function.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":concrete_function_type",
":function_metadata",
":function_metadata_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_status_internal",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:tfe_op_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/c/experimental/saved_model/core:concrete_function",
"//tensorflow/c/experimental/saved_model/core:function_metadata",
"//tensorflow/core:lib",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "concrete_function_list",
srcs = [
"concrete_function_list.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:concrete_function_list.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":concrete_function",
":concrete_function_list_type",
":concrete_function_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c/experimental/saved_model/core:concrete_function",
],
)
cc_library(
name = "concrete_function_list_type",
hdrs = [
"concrete_function_list_type.h",
],
deps = [
"//tensorflow/c/experimental/saved_model/core:concrete_function",
],
)
cc_library(
name = "concrete_function_type",
hdrs = [
"concrete_function_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:concrete_function",
],
)
cc_library(
name = "function_metadata",
srcs = [
"function_metadata.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:function_metadata.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":function_metadata_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c/experimental/saved_model/core:function_metadata",
],
)
cc_library(
name = "function_metadata_type",
hdrs = [
"function_metadata_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:function_metadata",
],
)
cc_library(
name = "saved_model_api",
srcs = [
"saved_model_api.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:saved_model_api.h",
],
copts = tf_copts(),
visibility = ["//tensorflow/c/experimental/saved_model/public:__pkg__"],
deps = [
":concrete_function",
":concrete_function_list",
":concrete_function_list_type",
":concrete_function_type",
":saved_model_api_type",
":signature_def_function",
":signature_def_function_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_status",
"//tensorflow/c:tf_status_internal",
"//tensorflow/c/eager:tfe_context_internal",
"//tensorflow/c/experimental/saved_model/core:saved_model_api",
"//tensorflow/c/experimental/saved_model/core:tf_saved_model_api",
"//tensorflow/core:lib",
"//tensorflow/core/common_runtime/eager:context",
"@com_google_absl//absl/base",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
],
)
cc_library(
name = "saved_model_api_type",
hdrs = [
"saved_model_api_type.h",
],
visibility = ["//visibility:private"],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:saved_model_api",
],
)
cc_library(
name = "signature_def_function",
srcs = [
"signature_def_function.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:signature_def_function.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":signature_def_function_metadata",
":signature_def_function_metadata_type",
":signature_def_function_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_status_internal",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:immediate_execution_operation",
"//tensorflow/c/eager:tfe_op_internal",
"//tensorflow/c/eager:tfe_tensorhandle_internal",
"//tensorflow/c/experimental/saved_model/core:signature_def_function",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
"//tensorflow/core:lib",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "signature_def_function_type",
hdrs = [
"signature_def_function_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:signature_def_function",
],
)
cc_library(
name = "signature_def_function_metadata",
srcs = [
"signature_def_function_metadata.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:signature_def_function_metadata.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":signature_def_function_metadata_type",
":signature_def_param_list",
":signature_def_param_list_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
],
)
cc_library(
name = "signature_def_function_metadata_type",
hdrs = [
"signature_def_function_metadata_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
],
)
cc_library(
name = "signature_def_param",
srcs = [
"signature_def_param.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:signature_def_param.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":signature_def_param_type",
":tensor_spec",
":tensor_spec_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_shape_internal",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
],
)
cc_library(
name = "signature_def_param_type",
hdrs = [
"signature_def_param_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
],
)
cc_library(
name = "signature_def_param_list",
srcs = [
"signature_def_param_list.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:signature_def_param_list.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":signature_def_param",
":signature_def_param_list_type",
":signature_def_param_type",
"//tensorflow/c:c_api_macros",
],
)
cc_library(
name = "signature_def_param_list_type",
hdrs = [
"signature_def_param_list_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
],
)
cc_library(
name = "tensor_spec",
srcs = [
"tensor_spec.cc",
],
hdrs = [
"//tensorflow/c/experimental/saved_model/public:tensor_spec.h",
],
copts = tf_copts(),
visibility = [
"//tensorflow/c/experimental/saved_model/public:__pkg__",
],
deps = [
":tensor_spec_type",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_datatype",
"//tensorflow/c:tf_shape",
"//tensorflow/c:tf_shape_internal",
"//tensorflow/c/experimental/saved_model/core:tensor_spec",
],
)
cc_library(
name = "tensor_spec_type",
hdrs = [
"tensor_spec_type.h",
],
deps = [
"//tensorflow/c:conversion_macros",
"//tensorflow/c:tf_shape_internal",
"//tensorflow/c/experimental/saved_model/core:tensor_spec",
],
)
tf_cc_test(
name = "saved_model_api_test",
size = "medium",
srcs = [
"saved_model_api_test.cc",
],
data = [
"//tensorflow/c/experimental/saved_model/internal/testdata:saved_models",
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
],
tags = [
"no_windows", # b/190030638
],
deps = [
":saved_model_api_type",
"//tensorflow/c:tf_datatype",
"//tensorflow/c:tf_shape",
"//tensorflow/c:tf_status",
"//tensorflow/c:tf_tensor",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/c/eager:c_api_test_util",
"//tensorflow/c/experimental/saved_model/core:tf_saved_model_api",
"//tensorflow/c/experimental/saved_model/public:concrete_function",
"//tensorflow/c/experimental/saved_model/public:saved_model_api",
"//tensorflow/c/experimental/saved_model/public:signature_def_function",
"//tensorflow/c/experimental/saved_model/public:signature_def_function_metadata",
"//tensorflow/c/experimental/saved_model/public:signature_def_param",
"//tensorflow/c/experimental/saved_model/public:signature_def_param_list",
"//tensorflow/c/experimental/saved_model/public:tensor_spec",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"@com_google_absl//absl/base",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
],
)
tensorflow/c/experimental/saved_model/internal/concrete_function.cc
+54
View File
@@ -0,0 +1,54 @@
/* 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/c/experimental/saved_model/public/concrete_function.h"
#include <cstddef>
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/tfe_op_internal.h"
#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/internal/concrete_function_type.h"
#include "tensorflow/c/experimental/saved_model/internal/function_metadata_type.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/platform/status.h"
extern "C" {
TF_FunctionMetadata* TF_ConcreteFunctionGetMetadata(TF_ConcreteFunction* func) {
return tensorflow::wrap(const_cast<tensorflow::FunctionMetadata*>(
&tensorflow::unwrap(func)->GetFunctionMetadata()));
}
TFE_Op* TF_ConcreteFunctionMakeCallOp(TF_ConcreteFunction* func,
TFE_TensorHandle** inputs, int num_inputs,
TF_Status* status) {
tensorflow::ImmediateOpPtr call_op;
absl::Span<tensorflow::AbstractTensorHandle* const> input_span(
reinterpret_cast<tensorflow::AbstractTensorHandle**>(
tensorflow::unwrap(inputs)),
static_cast<size_t>(num_inputs));
status->status = tensorflow::unwrap(func)->MakeCallOp(input_span, &call_op);
if (!status->status.ok()) {
return nullptr;
}
return tensorflow::wrap(call_op.release());
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/concrete_function_list.cc
+37
View File
@@ -0,0 +1,37 @@
/* 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 <stddef.h>
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/internal/concrete_function_list_type.h"
#include "tensorflow/c/experimental/saved_model/internal/concrete_function_type.h"
extern "C" {
size_t TF_ConcreteFunctionListNumOutputs(TF_ConcreteFunctionList* list) {
return list->list.size();
}
TF_ConcreteFunction* TF_ConcreteFunctionListGet(TF_ConcreteFunctionList* list,
int i) {
return tensorflow::wrap(list->list[i]);
}
void TF_DeleteConcreteFunctionList(TF_ConcreteFunctionList* list) {
delete list;
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/concrete_function_list_type.h
+30
View File
@@ -0,0 +1,30 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_LIST_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_LIST_TYPE_H_
#include <vector>
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
// Internal structures used by the SavedModel C API. These are likely to change
// and should not be depended on.
struct TF_ConcreteFunctionList {
std::vector<tensorflow::ConcreteFunction*> list;
};
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_LIST_TYPE_H_
tensorflow/c/experimental/saved_model/internal/concrete_function_type.h
+36
View File
@@ -0,0 +1,36 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
// Internal structures used by the SavedModel C API. These are likely to change
// and should not be depended on.
// It doesn't make sense to wrap tensorflow::ConcreteFunction* in a separate
// struct, since the lifetime of the struct and the raw pointer it wraps would
// be different. Therefore TF_ConcreteFunction* = tensorflow::ConcreteFunction*.
typedef struct TF_ConcreteFunction TF_ConcreteFunction;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::ConcreteFunction, TF_ConcreteFunction)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_CONCRETE_FUNCTION_TYPE_H_
tensorflow/c/experimental/saved_model/internal/function_metadata.cc
+20
View File
@@ -0,0 +1,20 @@
/* 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/c/experimental/saved_model/public/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/internal/function_metadata_type.h"
// TODO(bmzhao): Add getter functions here as necessary.
tensorflow/c/experimental/saved_model/internal/function_metadata_type.h
+30
View File
@@ -0,0 +1,30 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_FUNCTION_METADATA_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_FUNCTION_METADATA_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/function_metadata.h"
typedef struct TF_FunctionMetadata TF_FunctionMetadata;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::FunctionMetadata, TF_FunctionMetadata)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_FUNCTION_METADATA_TYPE_H_
tensorflow/c/experimental/saved_model/internal/saved_model_api.cc
+126
View File
@@ -0,0 +1,126 @@
/* 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/c/experimental/saved_model/public/saved_model_api.h"
#include <memory>
#include <optional>
#include <string>
#include <unordered_set>
#include <utility>
#include "absl/base/casts.h"
#include "absl/status/status.h"
#include "tensorflow/c/eager/tfe_context_internal.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_api.h"
#include "tensorflow/c/experimental/saved_model/core/tf_saved_model_api.h"
#include "tensorflow/c/experimental/saved_model/internal/concrete_function_list_type.h"
#include "tensorflow/c/experimental/saved_model/internal/concrete_function_type.h"
#include "tensorflow/c/experimental/saved_model/internal/saved_model_api_type.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_function_type.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/common_runtime/eager/context.h"
#include "tensorflow/core/platform/casts.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
extern "C" {
TF_SavedModel* TF_LoadSavedModel(const char* dirname, TFE_Context* ctx,
TF_Status* status) {
std::string saved_model_dir(dirname);
std::unique_ptr<tensorflow::SavedModelAPI> result;
if (tensorflow::unwrap(ctx)->UsesTFRT()) {
status->status = absl::UnimplementedError(
"TFRT SavedModel implementation will be added in the future");
} else {
std::unique_ptr<tensorflow::TFSavedModelAPI> saved_model;
status->status = tensorflow::TFSavedModelAPI::Load(
dirname, std::nullopt,
absl::down_cast<tensorflow::EagerContext*>(tensorflow::unwrap(ctx)),
&saved_model);
result = std::move(saved_model);
}
if (!status->status.ok()) {
return nullptr;
}
return tensorflow::wrap(result.release());
}
TF_SavedModel* TF_LoadSavedModelWithTags(const char* dirname, TFE_Context* ctx,
const char* const* tags, int tags_len,
TF_Status* status) {
std::string saved_model_dir(dirname);
std::unordered_set<std::string> tagset;
for (int i = 0; i < tags_len; ++i) {
tagset.insert(std::string(tags[i]));
}
std::unique_ptr<tensorflow::SavedModelAPI> result;
if (tensorflow::unwrap(ctx)->UsesTFRT()) {
status->status = absl::UnimplementedError(
"TFRT SavedModel implementation will be added in the future");
} else {
std::unique_ptr<tensorflow::TFSavedModelAPI> saved_model;
status->status = tensorflow::TFSavedModelAPI::Load(
dirname, tagset,
absl::down_cast<tensorflow::EagerContext*>(tensorflow::unwrap(ctx)),
&saved_model);
result = std::move(saved_model);
}
if (!status->status.ok()) {
return nullptr;
}
return tensorflow::wrap(result.release());
}
void TF_DeleteSavedModel(TF_SavedModel* model) {
delete tensorflow::unwrap(model);
}
TF_ConcreteFunction* TF_GetSavedModelConcreteFunction(TF_SavedModel* model,
const char* function_path,
TF_Status* status) {
tensorflow::ConcreteFunction* result = nullptr;
absl::Status get_function_status =
tensorflow::unwrap(model)->GetFunction(function_path, &result);
status->status.Update(get_function_status);
if (!get_function_status.ok()) {
return nullptr;
}
return tensorflow::wrap(result);
}
TF_CAPI_EXPORT extern TF_SignatureDefFunction*
TF_GetSavedModelSignatureDefFunction(TF_SavedModel* model,
const char* signature_def_key,
TF_Status* status) {
tensorflow::SignatureDefFunction* result = nullptr;
absl::Status get_function_status =
tensorflow::unwrap(model)->GetSignatureDefFunction(signature_def_key,
&result);
status->status.Update(get_function_status);
if (!get_function_status.ok()) {
return nullptr;
}
return tensorflow::wrap(result);
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/saved_model_api_test.cc
+587
View File
@@ -0,0 +1,587 @@
/* 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/c/experimental/saved_model/public/saved_model_api.h"
#include <cstdint>
#include <string>
#include <vector>
#include "absl/base/casts.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"
#include "tensorflow/c/eager/c_api_test_util.h"
#include "tensorflow/c/experimental/saved_model/core/tf_saved_model_api.h"
#include "tensorflow/c/experimental/saved_model/internal/saved_model_api_type.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param_list.h"
#include "tensorflow/c/experimental/saved_model/public/tensor_spec.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_shape.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/tstring.h"
namespace {
using tensorflow::tstring;
constexpr char kTestData[] = "cc/saved_model/testdata";
const char* kServeTag[] = {"serve"};
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 CSavedModelAPITest : public ::testing::TestWithParam<bool> {};
TEST_P(CSavedModelAPITest, LoadsSavedModelWithTags) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
TF_SavedModel* saved_model =
TF_LoadSavedModelWithTags(model_dir.c_str(), ctx, kServeTag, 1, status);
// 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(TF_GetCode(status), TF_UNIMPLEMENTED);
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
TEST_P(CSavedModelAPITest, LoadsSavedModel) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_ConcreteFunction* compute_fn =
TF_GetSavedModelConcreteFunction(saved_model, "compute", status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
std::vector<TFE_TensorHandle*> compute_fn_inputs;
TFE_TensorHandle* input_a = TestScalarTensorHandle(ctx, 2.0f);
TFE_TensorHandle* input_b = TestScalarTensorHandle(ctx, 1.0f);
compute_fn_inputs.push_back(input_a);
compute_fn_inputs.push_back(input_b);
TFE_Op* compute_fn_op = TF_ConcreteFunctionMakeCallOp(
compute_fn, compute_fn_inputs.data(), compute_fn_inputs.size(), status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// TODO(bmzhao): Finish API on FunctionMetadata args, so we know how many
// inputs + outputs a function has.
TFE_TensorHandle* compute_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(compute_fn_op, &compute_fn_outputs[0], &num_retvals, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(compute_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
float output_value = *static_cast<float*>(TF_TensorData(result));
// (1 + 2) * (2 + 1) / 3 + 5 should be 8
EXPECT_FLOAT_EQ(output_value, 8.0);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(compute_fn_outputs[0]);
TFE_DeleteTensorHandle(input_a);
TFE_DeleteTensorHandle(input_b);
TFE_DeleteOp(compute_fn_op);
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
// This tests running the "serving_default" SignatureDefFunction from the
// VarsAndArithmeticObjectGraph savedmodel. Here's what the signature_defs
// protobuf in the metagraph looks like:
// signature_def: {
// key : "serving_default"
// value: {
// inputs: {
// key : "a"
// value: {
// name : "serving_default_a:0"
// dtype: DT_FLOAT
// tensor_shape: {
// }
// }
// }
// inputs: {
// key : "b"
// value: {
// name : "serving_default_b:0"
// dtype: DT_FLOAT
// tensor_shape: {
// }
// }
// }
// outputs: {
// key : "output_0"
// value: {
// name : "StatefulPartitionedCall:0"
// dtype: DT_FLOAT
// tensor_shape: {
// }
// }
// }
// method_name: "tensorflow/serving/predict"
// }
// }
TEST_P(CSavedModelAPITest, RunsSignatureDefFunction) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = SavedModelPath("VarsAndArithmeticObjectGraph");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_SignatureDefFunction* serving_default =
TF_GetSavedModelSignatureDefFunction(saved_model, "serving_default",
status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_SignatureDefFunctionMetadata* metadata =
TF_SignatureDefFunctionGetMetadata(serving_default);
const TF_SignatureDefParamList* args =
TF_SignatureDefFunctionMetadataArgs(metadata);
const TF_SignatureDefParamList* returns =
TF_SignatureDefFunctionMetadataReturns(metadata);
EXPECT_EQ(TF_SignatureDefParamListSize(args), 2);
const TF_SignatureDefParam* param_a = TF_SignatureDefParamListGet(args, 0);
const TF_TensorSpec* tensor_spec_a = TF_SignatureDefParamTensorSpec(param_a);
const TF_Shape* shape_a = TF_TensorSpecShape(tensor_spec_a);
// Input "a" is a scalar, float32 tensor
EXPECT_EQ("a", std::string(TF_SignatureDefParamName(param_a)));
EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_a));
EXPECT_EQ(0, TF_ShapeDims(shape_a));
const TF_SignatureDefParam* param_b = TF_SignatureDefParamListGet(args, 1);
const TF_TensorSpec* tensor_spec_b = TF_SignatureDefParamTensorSpec(param_b);
const TF_Shape* shape_b = TF_TensorSpecShape(tensor_spec_b);
// Input "b" is a scalar, float32 tensor
EXPECT_EQ("b", std::string(TF_SignatureDefParamName(param_b)));
EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_b));
EXPECT_EQ(0, TF_ShapeDims(shape_b));
EXPECT_EQ(TF_SignatureDefParamListSize(returns), 1);
const TF_SignatureDefParam* param_out =
TF_SignatureDefParamListGet(returns, 0);
const TF_TensorSpec* tensor_spec_out =
TF_SignatureDefParamTensorSpec(param_out);
const TF_Shape* shape_out = TF_TensorSpecShape(tensor_spec_out);
// Output "output_0" is a scalar, float32 tensor
EXPECT_EQ("output_0", std::string(TF_SignatureDefParamName(param_out)));
EXPECT_EQ(TF_FLOAT, TF_TensorSpecDataType(tensor_spec_out));
EXPECT_EQ(0, TF_ShapeDims(shape_out));
std::vector<TFE_TensorHandle*> compute_fn_inputs;
TFE_TensorHandle* input_a = TestScalarTensorHandle(ctx, 2.0f);
TFE_TensorHandle* input_b = TestScalarTensorHandle(ctx, 1.0f);
compute_fn_inputs.push_back(input_a);
compute_fn_inputs.push_back(input_b);
TFE_Op* serving_default_op = TF_SignatureDefFunctionMakeCallOp(
serving_default, compute_fn_inputs.data(), compute_fn_inputs.size(),
status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
std::vector<TFE_TensorHandle*> compute_fn_outputs(
TF_SignatureDefParamListSize(returns));
int num_retvals = TF_SignatureDefParamListSize(returns);
TFE_Execute(serving_default_op, compute_fn_outputs.data(), &num_retvals,
status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(compute_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
float output_value = *static_cast<float*>(TF_TensorData(result));
// (1 + 2) * (2 + 1) / 3 + 5 should be 8
EXPECT_FLOAT_EQ(output_value, 8.0);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(compute_fn_outputs[0]);
TFE_DeleteTensorHandle(input_a);
TFE_DeleteTensorHandle(input_b);
TFE_DeleteOp(serving_default_op);
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
TEST_P(CSavedModelAPITest, LoadsAssetSavedModel) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = SavedModelPath("AssetModule");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_ConcreteFunction* read_file_fn =
TF_GetSavedModelConcreteFunction(saved_model, "read_file", status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TFE_Op* read_file_op =
TF_ConcreteFunctionMakeCallOp(read_file_fn, nullptr, 0, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// TODO(bmzhao): Finish API on FunctionMetadata args, so we know how many
// inputs + outputs a function has.
TFE_TensorHandle* read_file_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(read_file_op, &read_file_fn_outputs[0], &num_retvals, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(read_file_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
tensorflow::tstring* output_value =
static_cast<tensorflow::tstring*>(TF_TensorData(result));
std::string file_contents(*output_value);
EXPECT_NE(file_contents.find("TEST ASSET FILE CONTENTS"), std::string::npos);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(read_file_fn_outputs[0]);
TFE_DeleteOp(read_file_op);
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
TEST_P(CSavedModelAPITest, LoadsStaticHashtableSavedModel) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = SavedModelPath("StaticHashTableModule");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_ConcreteFunction* lookup_fn =
TF_GetSavedModelConcreteFunction(saved_model, "lookup", status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// Note(bmzhao): Based on static_hashtable_asset.txt, we expect the following
// mapping:
// "foo" -> 0
// "bar" -> 1
// "baz" -> 2
// "wombat" -> 3
// all other strings -> -1
// Call lookup function with input "foo", expecting an output of 0
{
std::vector<TFE_TensorHandle*> lookup_fn_inputs;
TFE_TensorHandle* input_foo = TestScalarTensorHandle(ctx, tstring("foo"));
lookup_fn_inputs.push_back(input_foo);
TFE_Op* lookup_op = TF_ConcreteFunctionMakeCallOp(
lookup_fn, lookup_fn_inputs.data(), lookup_fn_inputs.size(), status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// TODO(bmzhao): Finish API on FunctionMetadata args, so we know how many
// inputs + outputs a function has.
TFE_TensorHandle* lookup_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(lookup_op, &lookup_fn_outputs[0], &num_retvals, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(lookup_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
int64_t* output_value = static_cast<int64_t*>(TF_TensorData(result));
EXPECT_EQ(*output_value, 0);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(input_foo);
TFE_DeleteTensorHandle(lookup_fn_outputs[0]);
TFE_DeleteOp(lookup_op);
}
// Call lookup function with input "baz", expecting an output of 2
{
std::vector<TFE_TensorHandle*> lookup_fn_inputs;
TFE_TensorHandle* input_foo = TestScalarTensorHandle(ctx, tstring("baz"));
lookup_fn_inputs.push_back(input_foo);
TFE_Op* lookup_op = TF_ConcreteFunctionMakeCallOp(
lookup_fn, lookup_fn_inputs.data(), lookup_fn_inputs.size(), status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// TODO(bmzhao): Finish API on FunctionMetadata args, so we know how many
// inputs + outputs a function has.
TFE_TensorHandle* lookup_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(lookup_op, &lookup_fn_outputs[0], &num_retvals, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(lookup_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
int64_t* output_value = static_cast<int64_t*>(TF_TensorData(result));
EXPECT_EQ(*output_value, 2);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(input_foo);
TFE_DeleteTensorHandle(lookup_fn_outputs[0]);
TFE_DeleteOp(lookup_op);
}
// Call lookup function w/input "NON-EXISTENT-KEY", expecting an output of -1
{
std::vector<TFE_TensorHandle*> lookup_fn_inputs;
TFE_TensorHandle* input_foo =
TestScalarTensorHandle(ctx, tstring("NON-EXISTENT-KEY"));
lookup_fn_inputs.push_back(input_foo);
TFE_Op* lookup_op = TF_ConcreteFunctionMakeCallOp(
lookup_fn, lookup_fn_inputs.data(), lookup_fn_inputs.size(), status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
// TODO(bmzhao): Finish API on FunctionMetadata args, so we know how many
// inputs + outputs a function has.
TFE_TensorHandle* lookup_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(lookup_op, &lookup_fn_outputs[0], &num_retvals, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(lookup_fn_outputs[0], status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
EXPECT_EQ(TF_NumDims(result), 0);
int64_t* output_value = static_cast<int64_t*>(TF_TensorData(result));
EXPECT_EQ(*output_value, -1);
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(input_foo);
TFE_DeleteTensorHandle(lookup_fn_outputs[0]);
TFE_DeleteOp(lookup_op);
}
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
TEST_P(CSavedModelAPITest, LoadSavedModelWithUninitializedVariable) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = tensorflow::io::JoinPath(
tensorflow::testing::TensorFlowSrcRoot(),
"c/experimental/saved_model/internal/testdata/UninitializedVariable");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
EXPECT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
tensorflow::TFSavedModelAPI* model_api =
absl::down_cast<tensorflow::TFSavedModelAPI*>(
tensorflow::unwrap(saved_model));
tensorflow::Variable* uninitialized_variable;
ASSERT_EQ(absl::OkStatus(), model_api->GetVariable("uninitialized_variable",
&uninitialized_variable));
ASSERT_EQ(tensorflow::DT_FLOAT, uninitialized_variable->dtype());
ASSERT_EQ(absl::OkStatus(),
model_api->GetVariable("sub_module.uninitialized_variable",
&uninitialized_variable));
ASSERT_EQ(tensorflow::DT_INT64, uninitialized_variable->dtype());
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
TEST_P(CSavedModelAPITest, LoadSavedModelWithWhileLoop) {
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* opts = TFE_NewContextOptions();
bool use_tfrt = GetParam();
if (use_tfrt) {
TFE_DeleteContextOptions(opts);
TF_DeleteStatus(status);
GTEST_SKIP(); // TODO(chky) : Enable this once TFRT is open sourced.
}
TFE_ContextOptionsSetTfrt(opts, use_tfrt);
TFE_Context* ctx = TFE_NewContext(opts, status);
ASSERT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
std::string model_dir = tensorflow::io::JoinPath(
tensorflow::testing::TensorFlowSrcRoot(),
"c/experimental/saved_model/internal/testdata/SimpleWhileLoop");
TF_SavedModel* saved_model =
TF_LoadSavedModel(model_dir.c_str(), ctx, status);
ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_ConcreteFunction* while_fn =
TF_GetSavedModelConcreteFunction(saved_model, "compute", status);
ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
std::vector<TFE_TensorHandle*> while_fn_inputs;
while_fn_inputs.push_back(TestScalarTensorHandle(ctx, 10.0f));
TFE_Op* while_fn_op = TF_ConcreteFunctionMakeCallOp(
while_fn, while_fn_inputs.data(), while_fn_inputs.size(), status);
ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TFE_TensorHandle* while_fn_outputs[1] = {nullptr};
int num_retvals = 1;
TFE_Execute(while_fn_op, &while_fn_outputs[0], &num_retvals, status);
ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
TF_Tensor* result = TFE_TensorHandleResolve(while_fn_outputs[0], status);
ASSERT_EQ(TF_GetCode(status), TF_OK) << TF_Message(status);
ASSERT_EQ(TF_NumDims(result), 0);
float output_value = *static_cast<float*>(TF_TensorData(result));
ASSERT_FLOAT_EQ(output_value, 55); // 10+9+...+1
TF_DeleteTensor(result);
TFE_DeleteTensorHandle(while_fn_outputs[0]);
TFE_DeleteOp(while_fn_op);
TFE_DeleteTensorHandle(while_fn_inputs[0]);
TF_DeleteSavedModel(saved_model);
TF_DeleteStatus(status);
TFE_DeleteContext(ctx);
}
INSTANTIATE_TEST_SUITE_P(RuntimeAgnosticSavedModelTests, CSavedModelAPITest,
::testing::Bool());
} // namespace
tensorflow/c/experimental/saved_model/internal/saved_model_api_type.h
+35
View File
@@ -0,0 +1,35 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SAVED_MODEL_API_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SAVED_MODEL_API_TYPE_H_
#include <memory>
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/saved_model_api.h"
// Internal structures used by the SavedModel C API. These are likely to change
// and should not be depended on.
typedef struct TF_SavedModel TF_SavedModel;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::SavedModelAPI, TF_SavedModel)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SAVED_MODEL_API_TYPE_H_
tensorflow/c/experimental/saved_model/internal/signature_def_function.cc
+55
View File
@@ -0,0 +1,55 @@
/* 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/c/experimental/saved_model/public/signature_def_function.h"
#include <cstddef>
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"
#include "tensorflow/c/eager/tfe_op_internal.h"
#include "tensorflow/c/eager/tfe_tensorhandle_internal.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_function_metadata_type.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_function_type.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/platform/status.h"
extern "C" {
TF_SignatureDefFunctionMetadata* TF_SignatureDefFunctionGetMetadata(
TF_SignatureDefFunction* func) {
return tensorflow::wrap(const_cast<tensorflow::SignatureDefFunctionMetadata*>(
&tensorflow::unwrap(func)->GetFunctionMetadata()));
}
TFE_Op* TF_SignatureDefFunctionMakeCallOp(TF_SignatureDefFunction* func,
TFE_TensorHandle** inputs,
int num_inputs, TF_Status* status) {
tensorflow::ImmediateOpPtr call_op;
absl::Span<tensorflow::AbstractTensorHandle* const> input_span(
reinterpret_cast<tensorflow::AbstractTensorHandle**>(
tensorflow::unwrap(inputs)),
static_cast<size_t>(num_inputs));
status->status = tensorflow::unwrap(func)->MakeCallOp(input_span, &call_op);
if (!status->status.ok()) {
return nullptr;
}
return tensorflow::wrap(call_op.release());
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/signature_def_function_metadata.cc
+33
View File
@@ -0,0 +1,33 @@
/* 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/c/experimental/saved_model/public/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_function_metadata_type.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_param_list_type.h"
extern "C" {
extern const TF_SignatureDefParamList* TF_SignatureDefFunctionMetadataArgs(
const TF_SignatureDefFunctionMetadata* list) {
return tensorflow::wrap(&tensorflow::unwrap(list)->arguments());
}
extern const TF_SignatureDefParamList* TF_SignatureDefFunctionMetadataReturns(
const TF_SignatureDefFunctionMetadata* list) {
return tensorflow::wrap(&tensorflow::unwrap(list)->returns());
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/signature_def_function_metadata_type.h
+31
View File
@@ -0,0 +1,31 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_METADATA_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_METADATA_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
typedef struct TF_SignatureDefFunctionMetadata TF_SignatureDefFunctionMetadata;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::SignatureDefFunctionMetadata,
TF_SignatureDefFunctionMetadata)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_METADATA_TYPE_H_
tensorflow/c/experimental/saved_model/internal/signature_def_function_type.h
+31
View File
@@ -0,0 +1,31 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
typedef struct TF_SignatureDefFunction TF_SignatureDefFunction;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::SignatureDefFunction,
TF_SignatureDefFunction)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_FUNCTION_TYPE_H_
tensorflow/c/experimental/saved_model/internal/signature_def_param.cc
+33
View File
@@ -0,0 +1,33 @@
/* 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/c/experimental/saved_model/public/signature_def_param.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_param_type.h"
#include "tensorflow/c/experimental/saved_model/internal/tensor_spec_type.h"
extern "C" {
extern const char* TF_SignatureDefParamName(const TF_SignatureDefParam* param) {
return tensorflow::unwrap(param)->name().c_str();
}
extern const TF_TensorSpec* TF_SignatureDefParamTensorSpec(
const TF_SignatureDefParam* param) {
return tensorflow::wrap(&tensorflow::unwrap(param)->spec());
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/signature_def_param_list.cc
+35
View File
@@ -0,0 +1,35 @@
/* 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/c/experimental/saved_model/public/signature_def_param_list.h"
#include <cstddef>
#include "tensorflow/c/experimental/saved_model/internal/signature_def_param_list_type.h"
#include "tensorflow/c/experimental/saved_model/internal/signature_def_param_type.h"
extern "C" {
extern size_t TF_SignatureDefParamListSize(
const TF_SignatureDefParamList* list) {
return tensorflow::unwrap(list)->size();
}
extern const TF_SignatureDefParam* TF_SignatureDefParamListGet(
const TF_SignatureDefParamList* list, int i) {
return tensorflow::wrap(&tensorflow::unwrap(list)->at(i));
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/signature_def_param_list_type.h
+33
View File
@@ -0,0 +1,33 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_LIST_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_LIST_TYPE_H_
#include <vector>
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
typedef struct TF_SignatureDefParamList TF_SignatureDefParamList;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(std::vector<SignatureDefParam>,
TF_SignatureDefParamList)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_LIST_TYPE_H_
tensorflow/c/experimental/saved_model/internal/signature_def_param_type.h
+30
View File
@@ -0,0 +1,30 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
typedef struct TF_SignatureDefParam TF_SignatureDefParam;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::SignatureDefParam, TF_SignatureDefParam)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_SIGNATURE_DEF_PARAM_TYPE_H_
tensorflow/c/experimental/saved_model/internal/tensor_spec.cc
+32
View File
@@ -0,0 +1,32 @@
/* 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/c/experimental/saved_model/public/tensor_spec.h"
#include "tensorflow/c/experimental/saved_model/core/tensor_spec.h"
#include "tensorflow/c/experimental/saved_model/internal/tensor_spec_type.h"
#include "tensorflow/c/tf_shape_internal.h"
extern "C" {
TF_DataType TF_TensorSpecDataType(const TF_TensorSpec* spec) {
return static_cast<TF_DataType>(tensorflow::unwrap(spec)->dtype());
}
const TF_Shape* TF_TensorSpecShape(const TF_TensorSpec* spec) {
return tensorflow::wrap(&tensorflow::unwrap(spec)->shape());
}
} // end extern "C"
tensorflow/c/experimental/saved_model/internal/tensor_spec_type.h
+30
View File
@@ -0,0 +1,30 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_TENSOR_SPEC_TYPE_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_TENSOR_SPEC_TYPE_H_
#include "tensorflow/c/conversion_macros.h"
#include "tensorflow/c/experimental/saved_model/core/tensor_spec.h"
typedef struct TF_TensorSpec TF_TensorSpec;
namespace tensorflow {
DEFINE_CONVERSION_FUNCTIONS(tensorflow::TensorSpec, TF_TensorSpec)
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_INTERNAL_TENSOR_SPEC_TYPE_H_
tensorflow/c/experimental/saved_model/internal/testdata/BUILD
Vendored
+40
View File
@@ -0,0 +1,40 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("//tensorflow:tensorflow.default.bzl", "filegroup")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# Run this binary manually, with an argument pointing to the testdata/
# directory, to generate the test files used by the filegroup rule below.
py_binary(
name = "gen_saved_models",
srcs = ["gen_saved_models.py"],
strict_deps = True,
deps = [
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/module",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/saved_model",
"@absl_py//absl:app",
],
)
# Files generated by the binary above.
filegroup(
name = "saved_models",
srcs = glob([
"SimpleWhileLoop/**",
"UninitializedVariable/**",
]),
visibility = [
"//tensorflow/c/experimental/saved_model/internal:__pkg__",
],
)
tensorflow/c/experimental/saved_model/internal/testdata/SimpleWhileLoop/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/SimpleWhileLoop/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/SimpleWhileLoop/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/UninitializedVariable/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/UninitializedVariable/variables/variables.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/UninitializedVariable/variables/variables.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/c/experimental/saved_model/internal/testdata/gen_saved_models.py
Vendored
+101
View File
@@ -0,0 +1,101 @@
# 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.
# ==============================================================================
"""Creates saved models used for testing.
This executable should be run with an argument pointing to the testdata/ folder
in this directory. It will re-generate the saved models that are used for
testing.
"""
import os
from absl import app
from tensorflow.python.compat import v2_compat
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.module import module
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.saved_model import saved_model
def _gen_uninitialized_variable(base_dir):
"""Generates a saved model with an uninitialized variable."""
class SubModule(module.Module):
"""A module with an UninitializedVariable."""
def __init__(self):
self.uninitialized_variable = resource_variable_ops.UninitializedVariable(
name="uninitialized_variable", dtype=dtypes.int64)
class Module(module.Module):
"""A module with an UninitializedVariable."""
def __init__(self):
super(Module, self).__init__()
self.sub_module = SubModule()
self.initialized_variable = variables.Variable(
1.0, name="initialized_variable")
# An UninitializedVariable with the same name as the variable in the
# SubModule, but with a different type.
self.uninitialized_variable = resource_variable_ops.UninitializedVariable(
name="uninitialized_variable", dtype=dtypes.float32)
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def compute(self, value):
return self.initialized_variable + value
to_save = Module()
saved_model.save(
to_save, export_dir=os.path.join(base_dir, "UninitializedVariable"))
def _gen_simple_while_loop(base_dir):
"""Generates a saved model with a while loop."""
class Module(module.Module):
"""A module with a while loop."""
@def_function.function(
input_signature=[tensor_spec.TensorSpec((), dtypes.float32)])
def compute(self, value):
acc, _ = while_loop.while_loop(
cond=lambda acc, i: i > 0,
body=lambda acc, i: (acc + i, i - 1),
loop_vars=(constant_op.constant(0.0), value))
return acc
to_save = Module()
saved_model.save(
to_save, export_dir=os.path.join(base_dir, "SimpleWhileLoop"))
def main(args):
if len(args) != 2:
raise app.UsageError("Expected one argument (base_dir).")
_, base_dir = args
_gen_uninitialized_variable(base_dir)
_gen_simple_while_loop(base_dir)
if __name__ == "__main__":
v2_compat.enable_v2_behavior()
app.run(main)
tensorflow/c/experimental/saved_model/public/BUILD
+101
View File
@@ -0,0 +1,101 @@
# Experimental SavedModel C APIs for TensorFlow.
# See RFC https://github.com/tensorflow/community/pull/207
# All headers are on the public surface of Tensorflow's C API.
# Once moved out of experimental, these will be stable.
# The idea behind a separate public/ directory is to make apparent
# which headers are part of TF's public interface (and which headers)
# are implementation details. This structure allows us to also perform future
# programmatic checks that all "public" headers only include other "public"
# headers.
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"],
)
# TODO(bmzhao): Remove these exports_files and rules, swap with cc_public_library instead.
# cc_public_library would allows us to separate the header dep graph from header+srcs dep graph.
exports_files(
[
"concrete_function.h",
"concrete_function_list.h",
"function_metadata.h",
"saved_model_api.h",
"signature_def_function.h",
"signature_def_function_metadata.h",
"signature_def_param.h",
"signature_def_param_list.h",
"tensor_spec.h",
],
visibility = ["//tensorflow/c/experimental/saved_model/internal:__pkg__"],
)
# The purpose of this header is to provide insulation against
# future changes where we rename/move a public header, without
# forcing all clients to change their "#includes".
cc_library(
name = "c_saved_model_api",
hdrs = ["c_saved_model_api.h"],
deps = [
":concrete_function",
":concrete_function_list",
":function_metadata",
":saved_model_api",
":signature_def_function",
":signature_def_function_metadata",
":signature_def_param",
":signature_def_param_list",
":tensor_spec",
],
)
alias(
name = "concrete_function",
actual = "//tensorflow/c/experimental/saved_model/internal:concrete_function",
)
alias(
name = "concrete_function_list",
actual = "//tensorflow/c/experimental/saved_model/internal:concrete_function_list",
)
alias(
name = "function_metadata",
actual = "//tensorflow/c/experimental/saved_model/internal:function_metadata",
)
alias(
name = "saved_model_api",
actual = "//tensorflow/c/experimental/saved_model/internal:saved_model_api",
)
alias(
name = "signature_def_function",
actual = "//tensorflow/c/experimental/saved_model/internal:signature_def_function",
)
alias(
name = "signature_def_function_metadata",
actual = "//tensorflow/c/experimental/saved_model/internal:signature_def_function_metadata",
)
alias(
name = "signature_def_param",
actual = "//tensorflow/c/experimental/saved_model/internal:signature_def_param",
)
alias(
name = "signature_def_param_list",
actual = "//tensorflow/c/experimental/saved_model/internal:signature_def_param_list",
)
alias(
name = "tensor_spec",
actual = "//tensorflow/c/experimental/saved_model/internal:tensor_spec",
)
tensorflow/c/experimental/saved_model/public/README.md
+28
View File
@@ -0,0 +1,28 @@
# TensorFlow Saved Model C API
## Small ConcreteFunction Example
The following example loads a saved model from `"/path/to/model"` and
executes a function `f` taking no arguments and returning one single
value (error checking is omitted for simplicity):
```c
TF_Status* status = TF_NewStatus();
TFE_ContextOptions* ctx_options = TFE_NewContextOptions();
TFE_Context* ctx = TFE_NewContext(ctx_options, status);
TF_SavedModel* saved_model = TF_LoadSavedModel("/path/to/model", ctx, status);
TF_ConcreteFunction* f = TF_GetSavedModelConcreteFunction(saved_model, "f", status);
TFE_Op* op = TF_ConcreteFunctionMakeCallOp(f, NULL, 0, status);
TFE_TensorHandle* output;
int nouts = 1;
TFE_Execute(op, &output, &nouts, status);
TFE_DeleteTensorHandle(output);
TFE_DeleteOp(op);
TFE_DeleteSavedModel(saved_model);
TFE_DeleteContext(ctx);
TFE_DeleteContextOptions(ctx_options);
TF_DeleteStatus(status);
```
tensorflow/c/experimental/saved_model/public/c_saved_model_api.h
+31
View File
@@ -0,0 +1,31 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_C_SAVED_MODEL_API_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_C_SAVED_MODEL_API_H_
// IWYU pragma: begin_exports
#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function_list.h"
#include "tensorflow/c/experimental/saved_model/public/function_metadata.h"
#include "tensorflow/c/experimental/saved_model/public/saved_model_api.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param_list.h"
#include "tensorflow/c/experimental/saved_model/public/tensor_spec.h"
// IWYU pragma: end_exports
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_C_SAVED_MODEL_API_H_
tensorflow/c/experimental/saved_model/public/concrete_function.h
+58
View File
@@ -0,0 +1,58 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_H_
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/experimental/saved_model/public/function_metadata.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that corresponds to a Function loaded from a SavedModel.
// TODO(bmzhao): Work together w/srbs@ to make sure this composes w/the
// C++ Unified Eager/Graph API's AbstractFunction
typedef struct TF_ConcreteFunction TF_ConcreteFunction;
// Returns FunctionMetadata associated with `func`. Metadata's lifetime is
// bound to `func`, which is bound to the TF_SavedModel it was loaded from.
TF_CAPI_EXPORT extern TF_FunctionMetadata* TF_ConcreteFunctionGetMetadata(
TF_ConcreteFunction* func);
// Returns a TFE_Op suitable for executing this function. Caller must provide
// all function inputs in `inputs`, and must not add any additional inputs on
// the returned op. (i.e. don't call TFE_OpAddInput or TFE_OpAddInputList).
// The caller is responsible for deleting the returned TFE_Op. If op
// construction fails, `status` will be non-OK and the returned pointer will be
// null.
// TODO(bmzhao): Remove this function in a subsequent change; Design + implement
// a Function Execution interface for ConcreteFunction that accepts a tagged
// union of types (tensorflow::Value). This effectively requires moving much of
// the implementation of function.py/def_function.py to C++, and exposing a
// high-level API here. A strawman for what this interface could look like:
// TF_Value* TF_ExecuteFunction(TFE_Context*, TF_ConcreteFunction*, TF_Value*
// inputs, int num_inputs, TF_Status* status);
TF_CAPI_EXPORT extern TFE_Op* TF_ConcreteFunctionMakeCallOp(
TF_ConcreteFunction* func, TFE_TensorHandle** inputs, int num_inputs,
TF_Status* status);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_H_
tensorflow/c/experimental/saved_model/public/concrete_function_list.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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
#include <stddef.h>
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that is acts like a list of TF_ConcreteFunction pointers.
typedef struct TF_ConcreteFunctionList TF_ConcreteFunctionList;
// Returns the size of `list`.
TF_CAPI_EXPORT extern size_t TF_ConcreteFunctionListSize(
TF_ConcreteFunctionList* list);
// Returns the `i`th TF_ConcreteFunction in the list.
TF_CAPI_EXPORT extern TF_ConcreteFunction* TF_ConcreteFunctionListGet(
TF_ConcreteFunctionList* list, int i);
// Deletes `list`.
TF_CAPI_EXPORT extern void TF_DeleteConcreteFunctionList(
TF_ConcreteFunctionList* list);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_CONCRETE_FUNCTION_LIST_H_
tensorflow/c/experimental/saved_model/public/function_metadata.h
+35
View File
@@ -0,0 +1,35 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_FUNCTION_METADATA_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_FUNCTION_METADATA_H_
#include "tensorflow/c/c_api_macros.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type used to store any metadata associated with a function.
typedef struct TF_FunctionMetadata TF_FunctionMetadata;
// TODO(bmzhao): Add getters for fields as we determine what metadata
// we want to expose.
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_FUNCTION_METADATA_H_
tensorflow/c/experimental/saved_model/public/saved_model_api.h
+107
View File
@@ -0,0 +1,107 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SAVED_MODEL_API_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SAVED_MODEL_API_H_
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/public/concrete_function_list.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function.h"
#include "tensorflow/c/tf_status.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type representing a Tensorflow "SavedModel"
// (https://www.tensorflow.org/guide/saved_model) that we always pass by pointer
// to achieve ABI stability.
typedef struct TF_SavedModel TF_SavedModel;
// Load a SavedModel from `dirname`. We expect the SavedModel to contain a
// single Metagraph (as for those exported from TF2's `tf.saved_model.save`).
//
// Params:
// dirname - A directory filepath that the SavedModel is at.
// ctx - A TFE_Context containing optional load/TF runtime options.
// `ctx` must outlive the returned TF_SavedModel pointer.
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr. Otherwise, returns a newly created
// TF_SavedModel instance. It must be deleted by calling TF_DeleteSavedModel.
TF_CAPI_EXPORT extern TF_SavedModel* TF_LoadSavedModel(const char* dirname,
TFE_Context* ctx,
TF_Status* status);
// Load a SavedModel from `dirname`.
//
// Params:
// dirname - A directory filepath that the SavedModel is at.
// ctx - A TFE_Context containing optional load/TF runtime options.
// `ctx` must outlive the returned TF_SavedModel pointer.
// tags - char* array of SavedModel tags. We will load the metagraph matching
// the tags.
// tags_len - number of elements in the `tags` array.
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr. Otherwise, returns a newly created
// TF_SavedModel instance. It must be deleted by calling TF_DeleteSavedModel.
TF_CAPI_EXPORT extern TF_SavedModel* TF_LoadSavedModelWithTags(
const char* dirname, TFE_Context* ctx, const char* const* tags,
int tags_len, TF_Status* status);
// Deletes a TF_SavedModel, and frees any resources owned by it.
TF_CAPI_EXPORT extern void TF_DeleteSavedModel(TF_SavedModel* model);
// Retrieve a function from the TF2 SavedModel via function path.
//
// Params:
// model - The TF2 SavedModel to load a function from.
// function_path - A string containing the path from the root saved python
// object to a tf.function method.
// TODO(bmzhao): Add a detailed example of this with a
// python tf.module before moving this out of experimental.
// status - Set to OK on success and an appropriate error on failure.
// Returns:
// If status is not OK, returns nullptr. Otherwise, returns a
// TF_ConcreteFunction instance. The lifetime of this instance is
// "conceptually" bound to `model`. Once `model` is deleted, all
// `TF_ConcreteFunctions` retrieved from it are invalid, and have been deleted.
TF_CAPI_EXPORT extern TF_ConcreteFunction* TF_GetSavedModelConcreteFunction(
TF_SavedModel* model, const char* function_path, TF_Status* status);
// Retrieve a function from the TF SavedModel via a SignatureDef key.
//
// Params:
// model - The SavedModel to load a function from.
// signature_def_key - The string key of the 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
// TF_SignatureDefFunction instance. Once `model` is deleted, all
// `TF_SignatureDefFunctions` retrieved from it are invalid, and have been
// deleted.
TF_CAPI_EXPORT extern TF_SignatureDefFunction*
TF_GetSavedModelSignatureDefFunction(TF_SavedModel* model,
const char* signature_def_key,
TF_Status* status);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SAVED_MODEL_API_H_
tensorflow/c/experimental/saved_model/public/signature_def_function.h
+50
View File
@@ -0,0 +1,50 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that corresponds to a SignatureDefFunction loaded from a
// SavedModel.
typedef struct TF_SignatureDefFunction TF_SignatureDefFunction;
// Returns FunctionMetadata associated with `func`. Metadata's lifetime is
// bound to `func`, which is bound to the TF_SavedModel it was loaded from.
TF_CAPI_EXPORT extern TF_SignatureDefFunctionMetadata*
TF_SignatureDefFunctionGetMetadata(TF_SignatureDefFunction* func);
// Returns a TFE_Op suitable for executing this function. Caller must provide
// all function inputs in `inputs`, and must not add any additional inputs on
// the returned op. (i.e. don't call TFE_OpAddInput or TFE_OpAddInputList).
// The caller is responsible for deleting the returned TFE_Op. If op
// construction fails, `status` will be non-OK and the returned pointer will be
// null.
TF_CAPI_EXPORT extern TFE_Op* TF_SignatureDefFunctionMakeCallOp(
TF_SignatureDefFunction* func, TFE_TensorHandle** inputs, int num_inputs,
TF_Status* status);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_H_
tensorflow/c/experimental/saved_model/public/signature_def_function_metadata.h
+46
View File
@@ -0,0 +1,46 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param_list.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that corresponds to a SignatureDefFunction loaded from a
// SavedModel.
typedef struct TF_SignatureDefFunctionMetadata TF_SignatureDefFunctionMetadata;
// Retrieves the arguments of the SignatureDefFunction. The caller is not
// responsible for freeing the returned pointer.
TF_CAPI_EXPORT extern const TF_SignatureDefParamList*
TF_SignatureDefFunctionMetadataArgs(
const TF_SignatureDefFunctionMetadata* list);
// Retrieves the returns of the SignatureDefFunction. The caller is not
// responsible for freeing the returned pointer.
TF_CAPI_EXPORT extern const TF_SignatureDefParamList*
TF_SignatureDefFunctionMetadataReturns(
const TF_SignatureDefFunctionMetadata* list);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_FUNCTION_METADATA_H_
tensorflow/c/experimental/saved_model/public/signature_def_param.h
+44
View File
@@ -0,0 +1,44 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_H_
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/saved_model/public/tensor_spec.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that containing metadata of an input/output of a
// TF_SignatureDefFunction loaded from a SavedModel.
typedef struct TF_SignatureDefParam TF_SignatureDefParam;
// Returns the name of the given parameter. The caller is not responsible for
// freeing the returned char*.
TF_CAPI_EXPORT extern const char* TF_SignatureDefParamName(
const TF_SignatureDefParam* param);
// Returns the TensorSpec associated with the given parameter. The caller is
// not reponsible for freeing the returned TF_TensorSpec*.
TF_CAPI_EXPORT extern const TF_TensorSpec* TF_SignatureDefParamTensorSpec(
const TF_SignatureDefParam* param);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_H_
tensorflow/c/experimental/saved_model/public/signature_def_param_list.h
+44
View File
@@ -0,0 +1,44 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_LIST_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_LIST_H_
#include <stddef.h>
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/saved_model/public/signature_def_param.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type that containing metadata of an input/output of a
// ConcreteFunction loaded from a SavedModel.
typedef struct TF_SignatureDefParamList TF_SignatureDefParamList;
// Returns the size of `list`.
TF_CAPI_EXPORT extern size_t TF_SignatureDefParamListSize(
const TF_SignatureDefParamList* list);
// Returns the `i`th TF_SignatureDefParam in the list.
TF_CAPI_EXPORT extern const TF_SignatureDefParam* TF_SignatureDefParamListGet(
const TF_SignatureDefParamList* list, int i);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_SIGNATURE_DEF_PARAM_LIST_H_
tensorflow/c/experimental/saved_model/public/tensor_spec.h
+46
View File
@@ -0,0 +1,46 @@
/* 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_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_TENSOR_SPEC_H_
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_TENSOR_SPEC_H_
#include <stddef.h>
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_shape.h"
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
// An opaque type corresponding to TensorSpec
typedef struct TF_TensorSpec TF_TensorSpec;
// Returns the dtype associated with the TensorSpec.
TF_CAPI_EXPORT extern TF_DataType TF_TensorSpecDataType(
const TF_TensorSpec* spec);
// Returns the shape associated with the TensorSpec. The returned Shape is not
// owned by the caller. Caller must not call TF_DeleteShape on the returned
// shape.
TF_CAPI_EXPORT extern const TF_Shape* TF_TensorSpecShape(
const TF_TensorSpec* spec);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_PUBLIC_TENSOR_SPEC_H_