From 00a5ef689b0080ae2a73fe81da4673985078ec5c Mon Sep 17 00:00:00 2001
From: Brian Zhao <bmzhao@google.com>
Date: Fri, 18 Sep 2020 14:09:38 -0700
Subject: [PATCH] Support savedmodels that reference assets.
PiperOrigin-RevId: 332523998
Change-Id: Ia0f7cd75f79020c1b385ee6c02323c9d431ff86d
---
tensorflow/c/eager/BUILD | 2 +
tensorflow/c/eager/c_api_test_util.cc | 17 ++++++
tensorflow/c/eager/c_api_test_util.h | 5 ++
.../c/experimental/saved_model/core/BUILD | 5 +-
.../saved_model/core/revived_types/BUILD | 19 ++++++
.../saved_model/core/revived_types/asset.cc | 49 +++++++++++++++
.../saved_model/core/revived_types/asset.h | 50 +++++++++++++++
.../saved_model/core/saved_model_utils.cc | 14 +++++
.../saved_model/core/saved_model_utils.h | 8 +++
.../saved_model/core/tf_saved_model_api.cc | 20 +++---
.../internal/saved_model_api_test.cc | 57 ++++++++++++++++++
tensorflow/cc/saved_model/BUILD | 13 ++++
.../AssetModule/assets/test_asset.txt | 1 +
.../testdata/AssetModule/saved_model.pb | Bin 0 -> 8248 bytes
.../variables/variables.data-00000-of-00001 | Bin 0 -> 38 bytes
.../AssetModule/variables/variables.index | Bin 0 -> 144 bytes
.../testdata/generate_saved_models.py | 15 +++++
.../cc/saved_model/testdata/test_asset.txt | 1 +
18 files changed, 266 insertions(+), 10 deletions(-)
create mode 100644 tensorflow/c/experimental/saved_model/core/revived_types/asset.cc
create mode 100644 tensorflow/c/experimental/saved_model/core/revived_types/asset.h
create mode 100644 tensorflow/cc/saved_model/testdata/AssetModule/assets/test_asset.txt
create mode 100644 tensorflow/cc/saved_model/testdata/AssetModule/saved_model.pb
create mode 100644 tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.data-00000-of-00001
create mode 100644 tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.index
create mode 100644 tensorflow/cc/saved_model/testdata/test_asset.txt
diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD
index d259b32f339..da18424f936 100644
--- a/tensorflow/c/eager/BUILD
+++ b/tensorflow/c/eager/BUILD
@@ -619,6 +619,8 @@ tf_cuda_library(
":c_api",
":c_api_experimental",
"//tensorflow/c:c_test_util",
+ "//tensorflow/c:tf_datatype",
+ "//tensorflow/c:tf_tensor",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
diff --git a/tensorflow/c/eager/c_api_test_util.cc b/tensorflow/c/eager/c_api_test_util.cc
index fd68866f502..6eb5b521c50 100644
--- a/tensorflow/c/eager/c_api_test_util.cc
+++ b/tensorflow/c/eager/c_api_test_util.cc
@@ -17,12 +17,16 @@ limitations under the License.
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_experimental.h"
+#include "tensorflow/c/tf_datatype.h"
+#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/protobuf/cluster.pb.h"
using tensorflow::string;
+using tensorflow::tstring;
TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, float value) {
float data[] = {value};
@@ -36,6 +40,19 @@ TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, float value) {
return th;
}
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx,
+ const tensorflow::tstring& value) {
+ TF_Status* status = TF_NewStatus();
+ TF_Tensor* t = TFE_AllocateHostTensor(ctx, TF_STRING, nullptr, 0, status);
+ tstring* data = static_cast<tstring*>(TF_TensorData(t));
+ *data = value;
+ TFE_TensorHandle* th = TFE_NewTensorHandleFromTensor(ctx, t, status);
+ CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
+ TF_DeleteTensor(t);
+ TF_DeleteStatus(status);
+ return th;
+}
+
TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, int value) {
int data[] = {value};
TF_Status* status = TF_NewStatus();
diff --git a/tensorflow/c/eager/c_api_test_util.h b/tensorflow/c/eager/c_api_test_util.h
index 2f77ae5cf44..ad0c7c6340f 100644
--- a/tensorflow/c/eager/c_api_test_util.h
+++ b/tensorflow/c/eager/c_api_test_util.h
@@ -16,6 +16,7 @@ limitations under the License.
#define TENSORFLOW_C_EAGER_C_API_TEST_UTIL_H_
#include "tensorflow/c/eager/c_api.h"
+#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/tensorflow_server.pb.h"
@@ -28,6 +29,10 @@ TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, int value);
// Return a tensor handle containing a bool scalar
TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx, bool value);
+// Return a tensor handle containing a tstring scalar
+TFE_TensorHandle* TestScalarTensorHandle(TFE_Context* ctx,
+ const tensorflow::tstring& value);
+
// Return a tensor handle containing a 2x2 matrix of doubles
TFE_TensorHandle* DoubleTestMatrixTensorHandle(TFE_Context* ctx);
diff --git a/tensorflow/c/experimental/saved_model/core/BUILD b/tensorflow/c/experimental/saved_model/core/BUILD
index 2feb7c1b33e..4127fdfe0ee 100644
--- a/tensorflow/c/experimental/saved_model/core/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/BUILD
@@ -62,6 +62,7 @@ cc_library(
":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:tf_concrete_function",
"//tensorflow/c/experimental/saved_model/core/revived_types:variable",
@@ -69,6 +70,7 @@ cc_library(
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
+ "@com_google_absl//absl/types:span",
],
)
@@ -138,7 +140,6 @@ cc_library(
":saved_model_api",
":saved_model_utils",
":signature_def_function",
- "//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/saved_model/core/ops:restore_ops",
@@ -148,10 +149,10 @@ cc_library(
"//tensorflow/c/experimental/saved_model/core/revived_types:variable",
"//tensorflow/cc/saved_model:bundle_v2",
"//tensorflow/cc/saved_model:constants",
+ "//tensorflow/cc/saved_model:loader_util",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
- "//tensorflow/core/common_runtime/eager:tensor_handle",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:optional",
diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/BUILD b/tensorflow/c/experimental/saved_model/core/revived_types/BUILD
index 25cac39daa0..93a485e717f 100644
--- a/tensorflow/c/experimental/saved_model/core/revived_types/BUILD
+++ b/tensorflow/c/experimental/saved_model/core/revived_types/BUILD
@@ -8,6 +8,25 @@ package(
licenses = ["notice"], # Apache 2.0
)
+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",
+ ],
+)
+
cc_library(
name = "constant",
srcs = [
diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/asset.cc b/tensorflow/c/experimental/saved_model/core/revived_types/asset.cc
new file mode 100644
index 00000000000..5cc14d615f5
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/revived_types/asset.cc
@@ -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.
+==============================================================================*/
+
+#include "tensorflow/c/experimental/saved_model/core/revived_types/asset.h"
+
+#include <string>
+
+#include "tensorflow/c/eager/immediate_execution_context.h"
+#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
+#include "tensorflow/c/tensor_interface.h"
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/path.h"
+
+namespace tensorflow {
+
+Asset::Asset(ImmediateTensorHandlePtr handle)
+ : TensorHandleConvertible(std::move(handle)) {}
+
+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 errors::Internal(
+ "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 Status();
+}
+
+} // namespace tensorflow
diff --git a/tensorflow/c/experimental/saved_model/core/revived_types/asset.h b/tensorflow/c/experimental/saved_model/core/revived_types/asset.h
new file mode 100644
index 00000000000..c98bd9b5628
--- /dev/null
+++ b/tensorflow/c/experimental/saved_model/core/revived_types/asset.h
@@ -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_REVIVED_TYPES_ASSET_H_
+#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_REVIVED_TYPES_ASSET_H_
+
+#include <string>
+
+#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"
+
+namespace tensorflow {
+
+class Asset : public TensorHandleConvertible {
+ public:
+ static 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_
diff --git a/tensorflow/c/experimental/saved_model/core/saved_model_utils.cc b/tensorflow/c/experimental/saved_model/core/saved_model_utils.cc
index e79fd8d7001..8d1e0966ff7 100644
--- a/tensorflow/c/experimental/saved_model/core/saved_model_utils.cc
+++ b/tensorflow/c/experimental/saved_model/core/saved_model_utils.cc
@@ -100,6 +100,20 @@ Status ValidateSavedFunctionCompatibleWithFunctionDef(
} // namespace
+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 errors::FailedPrecondition(
+ "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);
+}
+
Status TensorProtoToConstant(ImmediateExecutionContext* ctx,
const TensorProto& proto,
std::unique_ptr<Constant>* output) {
diff --git a/tensorflow/c/experimental/saved_model/core/saved_model_utils.h b/tensorflow/c/experimental/saved_model/core/saved_model_utils.h
index 68bfbe32222..e82ec1bd104 100644
--- a/tensorflow/c/experimental/saved_model/core/saved_model_utils.h
+++ b/tensorflow/c/experimental/saved_model/core/saved_model_utils.h
@@ -22,7 +22,9 @@ limitations under the License.
#include <memory>
#include <unordered_map>
+#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/tf_concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/variable.h"
@@ -31,6 +33,7 @@ limitations under the License.
#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"
@@ -52,6 +55,11 @@ Status LoadSavedVariable(ImmediateExecutionContext* ctx,
const SavedVariable& variable,
std::unique_ptr<Variable>* output);
+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.
Status LoadTFConcreteFunction(
const SavedConcreteFunction& saved_concrete_function,
diff --git a/tensorflow/c/experimental/saved_model/core/tf_saved_model_api.cc b/tensorflow/c/experimental/saved_model/core/tf_saved_model_api.cc
index ab7052b52ed..143257b01d5 100644
--- a/tensorflow/c/experimental/saved_model/core/tf_saved_model_api.cc
+++ b/tensorflow/c/experimental/saved_model/core/tf_saved_model_api.cc
@@ -37,6 +37,7 @@ limitations under the License.
#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/cc/saved_model/loader_util.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
@@ -108,12 +109,17 @@ Status ConstantFromSavedConstant(
// SavedResources. These are returned via the `out` parameter.
Status ReviveObjects(
const MetaGraphDef& metagraph, ImmediateExecutionContext* context,
+ const std::string& directory,
std::unordered_map<int, std::unique_ptr<TensorHandleConvertible>>*
revived_objects) {
// This is needed to restore "Constant" nodes by looking up their
// "Value" attribute.
NodeAttrMap node_attr_map = internal::NodeToAttrMap(metagraph.graph_def());
+ // These are needed for creating "Assets", by looking up their filenames.
+ std::vector<AssetFileDef> assets;
+ TF_RETURN_IF_ERROR(internal::GetAssetFileDefs(metagraph, &assets));
+
// Iterate through all the saved objects, restoring objects as we go.
// We don't recreate functions until all other objects have been created.
for (int i = 0; i < metagraph.object_graph_def().nodes_size(); ++i) {
@@ -129,12 +135,10 @@ Status ReviveObjects(
node_attr_map, &constant));
(*revived_objects)[i] = std::move(constant);
} else if (node.kind_case() == SavedObject::kAsset) {
- // TODO(bmzhao): Implement Asset C++ class. This should be just recreating
- // the full path to the asset file:
- // https://github.com/tensorflow/tensorflow/blob/6a0bdbdb7c48a3491ae1277083ae3dafb4ab4d7a/tensorflow/python/saved_model/load.py#L395-L396
- // and storing it as a string tensor:
- // https://github.com/tensorflow/tensorflow/blob/6a0bdbdb7c48a3491ae1277083ae3dafb4ab4d7a/tensorflow/python/training/tracking/tracking.py#L324-L325
- return errors::Unimplemented("SavedAsset loading is not implemented yet");
+ std::unique_ptr<Asset> asset;
+ TF_RETURN_IF_ERROR(internal::LoadSavedAsset(context, node.asset(),
+ directory, assets, &asset));
+ (*revived_objects)[i] = std::move(asset);
} else if (node.kind_case() == SavedObject::kResource) {
// TODO(bmzhao): Figure out how resource loading works and implement it
return errors::Unimplemented(
@@ -352,8 +356,8 @@ Status TFSavedModelAPI::Load(
// https://github.com/tensorflow/tensorflow/blob/285b5fa15405c5e2c084080f52a1818be8648079/tensorflow/python/saved_model/function_deserialization.py#L438-L454
RevivedObjectMap revived_objects;
- TF_RETURN_IF_ERROR(
- ReviveObjects(bundle.meta_graph_def(), context, &revived_objects));
+ TF_RETURN_IF_ERROR(ReviveObjects(bundle.meta_graph_def(), context, directory,
+ &revived_objects));
// TODO(bmzhao): When we later add support for loading resources, we need to
// handle the case where materializing a function's captures requires invoking
diff --git a/tensorflow/c/experimental/saved_model/internal/saved_model_api_test.cc b/tensorflow/c/experimental/saved_model/internal/saved_model_api_test.cc
index df998fcf6cd..86754b32c0c 100644
--- a/tensorflow/c/experimental/saved_model/internal/saved_model_api_test.cc
+++ b/tensorflow/c/experimental/saved_model/internal/saved_model_api_test.cc
@@ -27,9 +27,12 @@ limitations under the License.
#include "tensorflow/core/lib/io/path.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"};
@@ -137,6 +140,60 @@ TEST_P(CSavedModelAPITest, LoadsSavedModel) {
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));
+ EXPECT_EQ(std::string(*output_value), "TEST ASSET FILE CONTENTS\n");
+
+ 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);
+}
+
INSTANTIATE_TEST_SUITE_P(RuntimeAgnosticSavedModelTests, CSavedModelAPITest,
::testing::Bool());
diff --git a/tensorflow/cc/saved_model/BUILD b/tensorflow/cc/saved_model/BUILD
index fddbcfec6e6..d25b580dace 100644
--- a/tensorflow/cc/saved_model/BUILD
+++ b/tensorflow/cc/saved_model/BUILD
@@ -209,9 +209,13 @@ tf_cc_test(
py_binary(
name = "testdata/generate_saved_models",
srcs = ["testdata/generate_saved_models.py"],
+ data = [
+ ":saved_model_asset_data",
+ ],
python_version = "PY3",
srcs_version = "PY3",
deps = [
+ "//tensorflow/python:client_testlib",
"//tensorflow/python:dtypes",
"//tensorflow/python:framework_ops",
"//tensorflow/python:tensor_spec",
@@ -221,6 +225,7 @@ py_binary(
"//tensorflow/python/module",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:save_options",
+ "//tensorflow/python/training/tracking",
"@absl_py//absl:app",
],
)
@@ -229,6 +234,7 @@ py_binary(
filegroup(
name = "saved_model_test_files",
srcs = glob([
+ "testdata/AssetModule/**",
"testdata/half_plus_two_pbtxt/**",
"testdata/half_plus_two_main_op/**",
"testdata/half_plus_two/**",
@@ -245,6 +251,13 @@ alias(
actual = ":saved_model_test_files",
)
+filegroup(
+ name = "saved_model_asset_data",
+ srcs = [
+ "testdata/test_asset.txt",
+ ],
+)
+
exports_files(
glob([
"testdata/half_plus_two_pbtxt/**",
diff --git a/tensorflow/cc/saved_model/testdata/AssetModule/assets/test_asset.txt b/tensorflow/cc/saved_model/testdata/AssetModule/assets/test_asset.txt
new file mode 100644
index 00000000000..40d69b1aac4
--- /dev/null
+++ b/tensorflow/cc/saved_model/testdata/AssetModule/assets/test_asset.txt
@@ -0,0 +1 @@
+TEST ASSET FILE CONTENTS
diff --git a/tensorflow/cc/saved_model/testdata/AssetModule/saved_model.pb b/tensorflow/cc/saved_model/testdata/AssetModule/saved_model.pb
new file mode 100644
index 0000000000000000000000000000000000000000..4bf99e03c22cd4c541b98d49e89ec15fccac46da
GIT binary patch
literal 8248
zcmdT}>u=lE6&LmPk!{&mekj*Y;xcLKXssnPFGtwgL~ha~T^z&qvZ2F5(BhTNMkWQ4
za*}-+R$#!8VaQNm1J(fp8VvnVV8F0qz_1VdL-uk1!LaYc&Lzc%DAIA3p<AF2B6%O@
z-1GdMbBhA`>KgoOocw(pS}=9pva~(@zG7<1ZkOHbk#Q)R%rg3Bn^kAXB;*d2Zl8%c
z)qc`r;sm%4B+FlaDhRL!xphOgY+8oA(YJej+nZY+tR;Dy>6T$e7Q+e@Z>vnVHTwwx
z$Z0xXE9}sczmc#bj)Ab_5(5V0E;IMoeQAB4wIB8jO}8yF50$q6D)&s*(H=9aS|D>k
zRA$+luGpHP4@Rm5aRJII>oS|kMpu<p&9q2%*D$&?y)5tp<srLa-0Q(z$Zjg_hbV%)
zBLWL~s|g#!?2Uu4K{Ck5vd|Dh%!*~mDvveG8KdcH3ceDcys4OYTI7bQ>q@svN;t{f
zv!u+9rc<Qj%1_Gf;DcBJ{+U2mv+QLq?eoimI0MdL<SPoPL$aN|-o^>}wxM^lJ(9O<
zQ`7fonxsSu<L*}!gJBy(V+%0B9<z4eHcXk52@k?Lls3Cco9&~lm^rHBVS7O8Ics0R
zJ={N<ZKkUNoP**PQ`DPUmvQsz;J>aMuxf$g8}fzbVWY<28cgzeNydm0dXX_2Ft&v{
zH%v4Vw5hyn4rmNhFzy<m%qes8n1LyWRawzh+3IR-{*_Kc!JTJGdB@@{u?mByVWBr1
zW9dSETRHp>sgakZ_8wA0({0;i0>)5qs$)!5AbAkD>+pYelndR0^7g)Bs?1Ra7X??y
zu3`-oLdF~gnrbzJb`N_0Kt>^%s#VQ!T?LwhzX~wFZ7VkG^t-Q}S#vkP1~Z(0t-Zze
z*yEf1ZuhQYxA(aOoJfSmfCl7xWEDhIRZp>PX6l}Qzy`p!<6?I(5}`)&&N{Rr9wmac
z1#vh?%)^+)dWxx_nR<J1Cxu@LP(s_kg$BtZF&B6ciA7w-T}9D!HndDU0VS>@raGn_
z-Z2tzMf4w>>3gSm2Fe(Pj7PXsWry4Y$hZ@2+{kgBrXj^ebrlHmb=$hFqrE9z?T1W_
zCJIVmQnO?&m{A2>D&{O^9<o|aTDiE=kP3bMp>90V*97wP9r!NHDVD`-83nr^T56JT
z%&E=6r0lAJ+hPxy%tE@Es+GJ{lx`IGi+eO2fgcM{^nwTpi?p@?dAZ$1s_1%UGSM~K
z&1>GwmnT$<h#rVf78E$|ZCoFcTs=mVG2*5ER!Mxp7L<JwZW<>e={(qmCJ1{-j_?Cm
z^!MbkxE=|}kpLZRpTy<x_NUz3N}LQ!rylsUT>jy3049iwDa47oYnaXT&F?fD$D&b_
zJlJu*^;WS!Qmr>&0mg~~$vM^rlf^Uv4tf;3kNyQI{M{SLJg<1rPtv2=KAQe%ghp-L
z&G<RUNJ61DF-Y};*OM?Ww~a23{@Njv4-{S5W2UUB7C|>n6wMM)*+fOkhTeUGPcE`U
zZ$q%<u6!qg4Nj5;C_n}!J|*yn5AY!+rSGKHginN2?L5aCXSYOHl5v*~Gnw9IvM*2h
zk*V~0$k^&7E>#PYqbYlN<5eB|m+9oBla!5Uga_)H7K+z_SBUncbm9M2+Gn}6Pj8)s
zg|M_`+eA!;umz<nUzFg}G@M9C#_}pTu4G;h`+JK^MBG|}Q*jbO1Lx>Zz9f;)m*AfS
z_K%btDMN`Hr(Yd#pva|qD50}qRf8x1aga7BQ{>fZ0khh3A*K+a&x8a-5sNcunC@_#
zt0g+~Yyd(q*F3N;H0$-1nlu+rn6{6KrpZ)>66>57O($9<Cz|F~5WCn?!7KZIjzMCb
ziwZ3*E@~f>&Nn?e7h0<hoo~oCJLt(*8q(SYwQ)u1T)29%v3q$<S-sTQT~#isS9YZf
z%5L4^X3Xb-{tPR)0Z1PF)01&PujyinW`j;T><8~tgX8#LqUtH|nv|oC1ZW}^Jr<}v
zgb_n{aY+BP>~ee*#SN0KrV}ZShw0Fd!q9!7m87NN1f8i7PYq+OLNP;9O+k=w@1PK%
zKZ4oq;S!>b$jek?m4B?19Y-}K%_q>IUtl2};ebRx&`N;JeWB7*k!3F9!FVM&rvGtR
z_lf6`5Kl&Sr@s}T;ueSXaP65OC^`aj(TbW`k}e8i6%}f$o{E-R=g?D}8LdTm{nm~3
z4>#}K-q^X``t}X^UhB?{^&R<xt?Qe&q}f<T=u=+9In^tdM303^;a6pH)?;$K1wNB0
z=|mzy`WNoE4l`u%PiaJ=O%F$)(=NtH)>j%`N;p<9V<^f;l{_<|3l6cc!R59DhzrRa
zH+MMZ%%2<WP1R!>Q^aX19eK3E1DE?KR~UA95D;p9{5isAq~+ssL_Zf|8*6bnu712j
z24YUma*bN7l%<(CW~dqJ75ZJn;TT$ZFv?tv2#=235YV9q0A7C*b6Y@l+WQ(DGmdtk
zrld51fz{heUKQU6;|x*>Mtm~0d6;mKyH^P#dTpkfDO7fhwn)sk=6W$j-h%{=;7yp3
z`&hEeR*$u{j^;HLCP-M3%|)!Hs)?KYEd$?radm*rPp=)|HUhN)$ilau9H)BnR=*yw
zb&12)g)Iz#Vc6<aW)1+o!ZS#ukw(8CSYycks*!1defWVPLBCSSyv37H^h8cZO}aK4
zwKb>f0>4SayNPEYvoOt%<ebP0)+5^eFw_qA8^d&TVG?5Dr{Hagn+V+L{a6uIyzwTY
zc&ELL-Gi?|I=?T#`!6g^qBTcMg+YHOz}%MC$^Ku%2Xys#Ug$3bpuy1B1R@rXqDlWC
zgzo@?>ww{8aXR2SjLm#kqGy5$(PR0OqS>W70vmXI<{5hgI+Q|8(yQUFSaj3)EIGKO
zqH3)9hl4fEJ>b__(QRMG8R+-Ii7FI{V#3A?_kHQ47axfJD~U|Z6`y6N#_1ir3ql|l
zG^&~nlGIQ^?3pwqudb!V(@^1isq%q=*NoVw({$PB5d`zDmTT|Atn74dyeYhC&ZoRd
z&Lj<iyB04Jh1w33c~c#)_ONT;QTkn*T*sBgL65wFH=@(g4AGb>&=jSvsdupDUmt8y
z#~U(LYunP>K%we1NJrZ7D2z!PW0C9(uh@_6=07rU>!_aU1lEpYesU8R5mxLmytvTX
zF1#>yKjybs*`k0w#&j(OZ$qV+CDDV4sQaL^c%bwKEB-7<f5O@&A(N6{3-C<%V+K0`
zC4R4S-Q3$oI72G#Z7V5#mO*WNoV8fDgV#2NL%fT@TUg#kMNooQjON2f&dVee2NHG%
zrvV;5;(b(SdgilC3U3i7-A_Xk36vCFH~94uc1B3%%S>uYsAZmKQgF6)jLs>!oP4+i
xA_#Qic?Q7QkMAuOo~E-Tg_oKi9Rqh=j@dx`GIJUb5coDU{WOiI_g``Ee*=T;X<q;U
literal 0
HcmV?d00001
diff --git a/tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.data-00000-of-00001 b/tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.data-00000-of-00001
new file mode 100644
index 0000000000000000000000000000000000000000..4105bb4c15e473b2f00dd3ae3236c7ef8c2328bf
GIT binary patch
literal 38
tcmY#<o|oOnCC$aj!6?L<SX`W1!o|nIB*aymnVy$eQd*Q+%*DXP006&g34Z_p
literal 0
HcmV?d00001
diff --git a/tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.index b/tensorflow/cc/saved_model/testdata/AssetModule/variables/variables.index
new file mode 100644
index 0000000000000000000000000000000000000000..3d903ca79a2904080c70501148466371b2e96f59
GIT binary patch
literal 144
zcmZQzVB=tvV&Y(Akl~JZ_HcFf4)FK%3vqPvagFzP@^W<!iFXfj4DjG!7h=#*GyNKU
yL4g4X7(vA3)xo)N&Vf%0H!v_VB`{dSg<QTJ^mO3x<pT*20^z?Kx>ZWuZvy~a=NKaZ
literal 0
HcmV?d00001
diff --git a/tensorflow/cc/saved_model/testdata/generate_saved_models.py b/tensorflow/cc/saved_model/testdata/generate_saved_models.py
index 5f39ae0651d..91c09d33bb0 100644
--- a/tensorflow/cc/saved_model/testdata/generate_saved_models.py
+++ b/tensorflow/cc/saved_model/testdata/generate_saved_models.py
@@ -29,9 +29,12 @@ from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_spec
from tensorflow.python.module import module
+from tensorflow.python.ops import io_ops
from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
from tensorflow.python.saved_model import save_options
from tensorflow.python.saved_model import saved_model
+from tensorflow.python.training.tracking import tracking
class VarsAndArithmeticObjectGraph(module.Module):
@@ -68,9 +71,21 @@ class CyclicModule(module.Module):
self.child = ReferencesParent(self)
+class AssetModule(module.Module):
+
+ def __init__(self):
+ self.asset = tracking.Asset(
+ test.test_src_dir_path("cc/saved_model/testdata/test_asset.txt"))
+
+ @def_function.function(input_signature=[])
+ def read_file(self):
+ return io_ops.read_file(self.asset)
+
+
MODULE_CTORS = {
"VarsAndArithmeticObjectGraph": VarsAndArithmeticObjectGraph,
"CyclicModule": CyclicModule,
+ "AssetModule": AssetModule,
}
diff --git a/tensorflow/cc/saved_model/testdata/test_asset.txt b/tensorflow/cc/saved_model/testdata/test_asset.txt
new file mode 100644
index 00000000000..40d69b1aac4
--- /dev/null
+++ b/tensorflow/cc/saved_model/testdata/test_asset.txt
@@ -0,0 +1 @@
+TEST ASSET FILE CONTENTS