Status InlineSymbolicGradient(const NodeDef& node,
FunctionOptimizerContext* ctx,
- GraphDef* inlined_graph) {
+ GraphDef* optimized_graph) {
VLOG(2) << "Inline symbolic gradient: " << SummarizeNodeDef(node);
GraphDef graph_def;
}
}
inlined_node.set_device(node.device());
- inlined_graph->add_node()->Swap(&inlined_node);
+ optimized_graph->add_node()->Swap(&inlined_node);
}
return Status::OK();
for (const NodeDef& node : item.graph.node()) {
const string func_name = node.op();
+ // Each node optimization can modify optimized graph only by adding new
+ // nodes, we can check node size to make sure that graph was not modified.
+ const int num_nodes_before = optimized_graph->node_size();
+ const auto is_graph_modified = [&]() {
+ int num_nodes = optimized_graph->node_size();
+ CHECK_GE(num_nodes, num_nodes_before) << "Nodes should not be removed";
+ return num_nodes > num_nodes_before;
+ };
+
+ // Add a copy of an input graph node to the optimized graph.
+ const auto add_node_copy = [&]() { *optimized_graph->add_node() = node; };
+
+// Skip errors if optimized graph was not modified before error happened.
+#define TF_SKIP_ERROR_IF_GRAPH_UNMODIFIED(...) \
+ do { \
+ const Status _status = (__VA_ARGS__); \
+ if (TF_PREDICT_FALSE(!_status.ok() && is_graph_modified())) \
+ return _status; \
+ if (TF_PREDICT_FALSE(!_status.ok() && !is_graph_modified())) { \
+ VLOG(3) << "Skip error: " << _status.error_message(); \
+ add_node_copy(); \
+ } \
+ } while (0)
+
+ // 1. Inline symbolic gradients into the optimized graph.
if (func_name == "SymbolicGradient" && inline_gradients) {
// Inline symbolic gradients only if the corresponding function is inlined
const auto* f_attr = gtl::FindOrNull(node.attr(), "f");
string f_name = f_attr != nullptr ? f_attr->func().name() : "";
if (ctx.IsInlinedFunction(f_name)) {
- TF_RETURN_IF_ERROR(InlineSymbolicGradient(node, &ctx, optimized_graph));
+ TF_SKIP_ERROR_IF_GRAPH_UNMODIFIED(
+ InlineSymbolicGradient(node, &ctx, optimized_graph));
continue;
}
}
+ // 2. Check if a node op is a function call.
const FunctionDef* func = ctx.function_library().Find(func_name);
if (func != nullptr) {
+ // 2a. Inline it if it's allowed to do so.
if (inline_func && ctx.IsInlinedFunction(func_name)) {
// Inline function body into the optimized graph}
- TF_RETURN_IF_ERROR(InlineFunction(node, *func, ctx, optimized_graph));
+ TF_SKIP_ERROR_IF_GRAPH_UNMODIFIED(
+ InlineFunction(node, *func, ctx, optimized_graph));
continue;
}
// Do not specialize if function has custom gradient.
const string grad_func = ctx.function_library().FindGradient(func_name);
+ // 2b. Specialize it to it's instantiation context if can't be inlined.
if (specialize_func && grad_func.empty() &&
(IsParametrized(*func) || HasTrulyConstInputs(node, ctx))) {
// TODO(ezhulenev): Specialize function call if input has a known shape.
// Specialize function body for its instantiation attributes and inputs.
- TF_RETURN_IF_ERROR(
+ TF_SKIP_ERROR_IF_GRAPH_UNMODIFIED(
SpecializeFunction(node, *func, &ctx, optimized_graph));
continue;
}
// If we reached this point, node was not handled by any of the stages
// (inline, specialize), simply add a copy to the graph.
- *optimized_graph->add_node() = node;
+ add_node_copy();
+
+#undef TF_SKIP_ERROR_IF_GRAPH_UNMODIFIED
}
*optimized_graph->mutable_versions() = item.graph.versions();
test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
}
+TEST_F(FunctionOptimizerTest, InlineFunction_SkipErrorsIfGraphNotModified) {
+ using test::function::NDef;
+
+ FunctionOptimizer optimizer(RewriterConfig::DEFAULT);
+
+ // Standard XTimesTwo() function.
+ FunctionDef x_times_two = test::function::XTimesTwo();
+
+ // Function with sequence of tensors as an input (currently not supported).
+ FunctionDef my_identity_n = FunctionDefHelper::Create(
+ // Name
+ "MyIdentityN",
+ // Args
+ {"x: N*T"},
+ // Return values
+ {"out: N*T"},
+ // Attrs
+ {"N:int", "T:{float, double, int32, int64}"},
+ // Nodes (just forward inputs through IdentityN)
+ {
+ {{"Id"}, "IdentityN", {"x"}, {{"T", "$T"}, {"N", "$N"}}},
+ },
+ // Output mapping
+ {{"out", "Id:output:0"}});
+
+ GrapplerItem item;
+ item.graph = test::function::GDef(
+ {NDef("x", "Placeholder", {}, {{"dtype", DT_FLOAT}}, kDevice),
+ NDef("y1", "XTimesTwo", {"x"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("y2", "MyIdentityN", {"x"}, {{"T", DT_FLOAT}, {"N", 1}}, kDevice),
+ NDef("z1", "Identity", {"y1:0"}, {{"T", DT_FLOAT}}, kDevice),
+ NDef("z2", "Identity", {"y2:0"}, {{"T", DT_FLOAT}}, kDevice)},
+ // FunctionLib
+ {x_times_two, my_identity_n});
+
+ GraphDef output;
+ TF_EXPECT_OK(optimizer.Optimize(nullptr, item, &output));
+
+ // Verify that only MyIdentityN is in the function library after optimization.
+ ASSERT_EQ(1, output.library().function().size());
+ EXPECT_EQ("MyIdentityN", output.library().function(0).signature().name());
+
+ // And that XTimesTwo was successfully inlined.
+ int found = 0;
+ for (const NodeDef& node : output.node()) {
+ if (node.name() == "y1/inlined_inputs") {
+ found++;
+ EXPECT_EQ("IdentityN", node.op());
+ EXPECT_EQ(kDevice, node.device());
+ EXPECT_EQ(1, node.input_size());
+ EXPECT_EQ("x", node.input(0));
+ } else if (node.name() == "y1") {
+ found++;
+ EXPECT_EQ("IdentityN", node.op());
+ EXPECT_EQ(kDevice, node.device());
+ EXPECT_EQ(1, node.input_size());
+ EXPECT_EQ("y1/y", node.input(0));
+ } else if (node.name() == "y2") {
+ found++;
+ EXPECT_EQ("MyIdentityN", node.op());
+ EXPECT_EQ(kDevice, node.device());
+ EXPECT_EQ(1, node.input_size());
+ EXPECT_EQ("x", node.input(0));
+ }
+ }
+ EXPECT_EQ(3, found);
+
+ Tensor pi = test::AsScalar<float>(3.14f);
+ item.fetch = {"z1"};
+ item.feed.emplace_back("x", pi);
+ auto tensors_expected = EvaluateFetchNodes(item);
+ GrapplerItem optimized(item, std::move(output));
+ auto tensors = EvaluateFetchNodes(optimized);
+ test::ExpectTensorEqual<float>(tensors_expected[0], tensors[0]);
+}
+
TEST_F(FunctionOptimizerTest, InlineFunction_FixedTypeFunction) {
using test::function::NDef;
projects / platform / upstream / tensorflow.git / commitdiff