ae07e29a16
* This fixes a correctness problem where we were relying on the variable's "name" attribute, which is merely advisory. * Should match the load heuristics in Python by reading the TrackableObjectGraph from the checkpoint, re-associating it with the SavedObjectGraph and using that to restore the variable. * Has a side-effect of eliminating the dependency of the importer on the CPU runtime and kernels, which should reduce necessary dependencies to compile by tens of megabytes. Open questions: * I created a new bundle_v2.h because V2 is so substantially different from V1. Also, it has a different dependency surface area and I therefore wanted to use it as a different (lighter) library. Advise if you would like this organized differently. I tried to factor it so that a future C++ SavedModel loader/runner could be built on this as well (with some incremental work). * There didn't seem to be a facility for generating the testdata SavedModels. I ended up just writing a small standalone python script that I can invoke from the command line to generate the new ones. It isn't wired up in any way but it should be possible to do so at some point (I'm not familiar enough with the test infra on this side of the tree, so please advise). * I could see factoring the SavedModelV2Bundle::RestoreObjects method into a dedicated facility for loading arbitrary checkpoints. Not sure what the roadmap here is so just kept is simple for now. PiperOrigin-RevId: 282439157 Change-Id: I685df47ab621917eed270319cde220498f191b74
100 lines
3.6 KiB
C++
100 lines
3.6 KiB
C++
/* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#include "tensorflow/cc/saved_model/bundle_v2.h"
|
|
|
|
#include "tensorflow/core/lib/core/status_test_util.h"
|
|
#include "tensorflow/core/lib/io/path.h"
|
|
#include "tensorflow/core/platform/test.h"
|
|
|
|
namespace tensorflow {
|
|
namespace {
|
|
|
|
constexpr char kTestData[] = "cc/saved_model/testdata";
|
|
|
|
class BundleV2Test : public ::testing::Test {
|
|
protected:
|
|
BundleV2Test() {}
|
|
|
|
void RestoreVarsAndVerify(SavedModelV2Bundle* bundle,
|
|
std::vector<std::string> expected_names) {
|
|
// Collect saved_node_id, full_name, checkpoint_key into a vector.
|
|
using RestoredVarType = std::tuple<int, std::string, std::string>;
|
|
std::vector<RestoredVarType> restored_vars;
|
|
TF_ASSERT_OK(bundle->VisitObjectsToRestore(
|
|
[&](int saved_node_id,
|
|
const TrackableObjectGraph::TrackableObject& trackable_object)
|
|
-> Status {
|
|
for (const auto& attr : trackable_object.attributes()) {
|
|
if (attr.name() == "VARIABLE_VALUE") {
|
|
restored_vars.emplace_back(saved_node_id, attr.full_name(),
|
|
attr.checkpoint_key());
|
|
}
|
|
}
|
|
return Status::OK();
|
|
}));
|
|
|
|
// Should be one of each var name restored.
|
|
for (const auto& expected_name : expected_names) {
|
|
EXPECT_EQ(1, std::count_if(restored_vars.begin(), restored_vars.end(),
|
|
[&](RestoredVarType t) {
|
|
return std::get<1>(t) == expected_name;
|
|
}));
|
|
}
|
|
|
|
for (const auto& restored_var : restored_vars) {
|
|
// Each restored var should match a SavedObjectGraph node with the same
|
|
// variable name.
|
|
const auto& saved_node =
|
|
bundle->saved_object_graph().nodes(std::get<0>(restored_var));
|
|
EXPECT_EQ(std::get<1>(restored_var), saved_node.variable().name());
|
|
|
|
// And should be able to load it from the tensor_bundle.
|
|
Tensor value;
|
|
TF_ASSERT_OK(
|
|
bundle->variable_reader()->Lookup(std::get<2>(restored_var), &value));
|
|
}
|
|
}
|
|
};
|
|
|
|
TEST_F(BundleV2Test, LoadsVarsAndArithmeticObjectGraph) {
|
|
const string export_dir = io::JoinPath(
|
|
testing::TensorFlowSrcRoot(), kTestData, "VarsAndArithmeticObjectGraph");
|
|
|
|
SavedModelV2Bundle bundle;
|
|
TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
|
|
|
|
// Ensure that there are nodes in the trackable_object_graph.
|
|
EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
|
|
|
|
RestoreVarsAndVerify(&bundle, {"variable_x", "variable_y", "child_variable"});
|
|
}
|
|
|
|
TEST_F(BundleV2Test, LoadsCyclicModule) {
|
|
const string export_dir =
|
|
io::JoinPath(testing::TensorFlowSrcRoot(), kTestData, "CyclicModule");
|
|
|
|
SavedModelV2Bundle bundle;
|
|
TF_ASSERT_OK(SavedModelV2Bundle::Load(export_dir, &bundle));
|
|
|
|
// Ensure that there are nodes in the trackable_object_graph.
|
|
EXPECT_GT(bundle.trackable_object_graph().nodes_size(), 0);
|
|
|
|
RestoreVarsAndVerify(&bundle, {"MyVariable"});
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace tensorflow
|