/* Copyright 2017 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/client/client_session.h"
|
#include "tensorflow/cc/framework/gradients.h"
|
#include "tensorflow/cc/framework/testutil.h"
|
#include "tensorflow/cc/ops/standard_ops.h"
|
#include "tensorflow/cc/ops/while_loop.h"
|
#include "tensorflow/core/framework/graph.pb.h"
|
#include "tensorflow/core/framework/tensor_testutil.h"
|
#include "tensorflow/core/lib/core/status_test_util.h"
|
#include "tensorflow/core/platform/test.h"
|
|
namespace tensorflow {
|
|
namespace {
|
|
class WhileGradientsTest : public ::testing::Test {
|
protected:
|
WhileGradientsTest() : scope_(Scope::NewRootScope()) {}
|
|
void Init(int num_inputs, DataType dtype = DT_INT32) {
|
for (int i = 0; i < num_inputs; ++i) {
|
inputs_.push_back(ops::Placeholder(scope_, dtype));
|
}
|
}
|
|
void CreateLoop(const ops::CondGraphBuilderFn& cond,
|
const ops::BodyGraphBuilderFn& body,
|
const std::vector<Output>* inputs = nullptr) {
|
if (inputs == nullptr) inputs = &inputs_;
|
TF_ASSERT_OK(ops::BuildWhileLoop(scope_, *inputs, cond, body, "test_loop",
|
&outputs_));
|
}
|
|
void CreateBackprop() {
|
TF_ASSERT_OK(
|
AddSymbolicGradients(scope_, outputs_, inputs_, &grad_outputs_));
|
ASSERT_EQ(grad_outputs_.size(), inputs_.size());
|
}
|
|
template <typename T>
|
void Run(const std::vector<Input::Initializer>& input_values,
|
const std::vector<T>& expected_grad_values) {
|
Run<T>(ClientSession(scope_), input_values, expected_grad_values);
|
}
|
|
template <typename T>
|
void Run(const ClientSession& session,
|
const std::vector<Input::Initializer>& input_values,
|
const std::vector<T>& expected_grad_values,
|
const RunOptions& run_options = RunOptions(),
|
RunMetadata* run_metadata = nullptr) {
|
DCHECK_EQ(input_values.size(), inputs_.size());
|
ClientSession::FeedType feeds;
|
for (int i = 0; i < inputs_.size(); ++i) {
|
feeds.emplace(inputs_[i], input_values[i]);
|
}
|
|
std::vector<Operation> run_outputs;
|
std::vector<Tensor> out_tensors;
|
TF_ASSERT_OK(session.Run(run_options, feeds, grad_outputs_, run_outputs,
|
&out_tensors, run_metadata));
|
ASSERT_EQ(out_tensors.size(), grad_outputs_.size());
|
|
DCHECK_EQ(expected_grad_values.size(), out_tensors.size());
|
for (int i = 0; i < out_tensors.size(); ++i) {
|
test::ExpectTensorEqual<T>(
|
out_tensors[i], test::AsTensor<T>({expected_grad_values[i]}, {}));
|
}
|
}
|
|
Scope scope_;
|
std::vector<Output> inputs_;
|
std::vector<Output> outputs_;
|
std::vector<Output> grad_outputs_;
|
};
|
|
TEST_F(WhileGradientsTest, Basic) {
|
// Create loop: while (i < 10) i += 1
|
Init(1);
|
CreateLoop(
|
[](const Scope& s, const std::vector<Output>& inputs, Output* output) {
|
*output = ops::Less(s, inputs[0], 10);
|
return s.status();
|
},
|
[](const Scope& s, const std::vector<Output>& inputs,
|
std::vector<Output>* outputs) {
|
// Use AddN, rather than Add, because the gradient function doesn't
|
// depend on the input shapes, and thus we do not need to store
|
// intermediate values in a stack.
|
outputs->push_back(ops::AddN(s, {inputs[0], 1}));
|
return s.status();
|
});
|
CreateBackprop();
|
|
Run<int>({1}, {1});
|
Run<int>({11}, {1});
|
}
|
|
TEST_F(WhileGradientsTest, MultipleLoopVars) {
|
// Create loop: while (i < 10) i += j; j += 1; k = k
|
Init(3);
|
CreateLoop(
|
[](const Scope& s, const std::vector<Output>& inputs, Output* output) {
|
*output = ops::Less(s, inputs[0], 10);
|
return s.status();
|
},
|
[](const Scope& s, const std::vector<Output>& inputs,
|
std::vector<Output>* outputs) {
|
outputs->push_back(ops::AddN(s, {inputs[0], inputs[1]}));
|
outputs->push_back(ops::AddN(s, {inputs[1], 1}));
|
outputs->push_back(inputs[2]);
|
return s.status();
|
});
|
CreateBackprop();
|
|
// The following execution traces illustrate why we expect dF/dj to be 5:
|
//
|
// i j k
|
// ---------
|
// 0 1 2 <-- initial values
|
// 1 2 2
|
// 3 3 2
|
// 6 4 2
|
// 10 5 2 <-- while output values
|
// outputs sum = 17
|
//
|
// i j k
|
// ---------
|
// 0 2 2 <-- initial values (add 1 to j)
|
// 2 3 2
|
// 5 4 2
|
// 9 5 2
|
// 14 6 2 <-- while output values
|
// outputs sum = 22
|
//
|
// Calculate the "slope" between j=1 and j=2:
|
// 22 - 17 = 5 => dF/dj = 5
|
Run<int>({0, 1, 2}, {1, 5, 1});
|
|
Run<int>({1, 1, 0}, {1, 5, 1});
|
Run<int>({0, 0, 0}, {1, 6, 1});
|
}
|
|
TEST_F(WhileGradientsTest, Chaining) {
|
Init(2, DT_DOUBLE);
|
|
// Multiply each input by 2 before passing to while loop to make sure chaining
|
// works properly
|
std::vector<Output> loop_inputs = {ops::Multiply(scope_, inputs_[0], 2.0),
|
ops::Multiply(scope_, inputs_[1], 2.0)};
|
|
// Create loop: while (i > 0 && j > 0) i -= 1
|
CreateLoop(
|
[](const Scope& s, const std::vector<Output>& inputs, Output* output) {
|
*output = ops::LogicalAnd(s, ops::Greater(s, inputs[0], 0.0),
|
ops::Greater(s, inputs[1], 0.0));
|
return s.status();
|
},
|
[](const Scope& s, const std::vector<Output>& inputs,
|
std::vector<Output>* outputs) {
|
outputs->push_back(ops::AddN(s, {inputs[0], -1.0}));
|
outputs->push_back(inputs[1]);
|
return s.status();
|
},
|
&loop_inputs);
|
|
// Take negative of first output to make sure chaining works properly
|
outputs_[0] = ops::Neg(scope_, outputs_[0]);
|
|
CreateBackprop();
|
|
Run<double>({1.0, 1.0}, {-2.0, 2.0});
|
Run<double>({0.0, 0.0}, {-2.0, 2.0});
|
}
|
|
TEST_F(WhileGradientsTest, MultipleDevices) {
|
// Make sure loop is created on cpu0
|
scope_ = scope_.WithDevice("/cpu:0");
|
|
// Create loop: while (i < 10) i += j
|
Init(2);
|
CreateLoop(
|
[](const Scope& s, const std::vector<Output>& inputs, Output* output) {
|
*output = ops::Less(s, inputs[0], 10);
|
return s.status();
|
},
|
[](const Scope& s, const std::vector<Output>& inputs,
|
std::vector<Output>* outputs) {
|
// Place body on cpu1
|
Scope cpu1_scope = s.WithDevice("/cpu:1");
|
outputs->push_back(ops::AddN(cpu1_scope, {inputs[0], inputs[1]}));
|
outputs->push_back(inputs[1]);
|
return cpu1_scope.status();
|
});
|
|
// Build gradient graph on cpu1
|
Scope cpu1_scope = scope_.WithDevice("/cpu:1");
|
TF_ASSERT_OK(
|
AddSymbolicGradients(cpu1_scope, outputs_, inputs_, &grad_outputs_));
|
ASSERT_EQ(grad_outputs_.size(), inputs_.size());
|
|
// Run with two CPU devices and output partition graphs
|
SessionOptions session_options;
|
(*session_options.config.mutable_device_count())["CPU"] = 2;
|
RunOptions run_options;
|
run_options.set_output_partition_graphs(true);
|
RunMetadata run_metadata;
|
Run<int>(ClientSession(scope_, session_options), {0, 1}, {1, 11}, run_options,
|
&run_metadata);
|
|
// Check that at least one node ran on each device
|
ASSERT_EQ(run_metadata.partition_graphs().size(), 2);
|
for (const GraphDef& partition_graph : run_metadata.partition_graphs()) {
|
EXPECT_GE(partition_graph.node().size(), 1);
|
}
|
}
|
|
} // namespace
|
} // namespace tensorflow
|
