Skip to content

Commit 2db9a1e

Browse files
tensorflower-gardenertensorflower-gardener
committed
C++ Gradients: Adds gradient functions and tests for Pack/Unpack, and takes care of multiple TODOS:
*) Adds support and unit test for returning dependent gradient outputs. *) Adds support and unit test for stopping backprop at frontier of requested inputs. *) Adds support and unit test for returning gradients for nodes with multiple outputs. *) Moves common unit test code out into a testlib. *) Moves common gradient-specific unit test code out into a separate testlib. Change: 132434513
1 parent d1e2c2d commit 2db9a1e

11 files changed

Lines changed: 584 additions & 123 deletions

File tree

‎tensorflow/cc/BUILD‎

Lines changed: 69 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ cc_library(
2121
deps = [
2222
":cc_ops",
2323
":grad_op_registry",
24-
":math_grad",
2524
":ops",
2625
":scope",
2726
"//tensorflow/core:core_cpu",
@@ -36,7 +35,9 @@ tf_cc_test(
3635
deps = [
3736
":cc_ops",
3837
":grad_op_registry",
38+
":grad_ops",
3939
":gradients",
40+
":testutil",
4041
"//tensorflow/core:all_kernels",
4142
"//tensorflow/core:core_cpu_internal",
4243
"//tensorflow/core:framework",
@@ -46,6 +47,26 @@ tf_cc_test(
4647
],
4748
)
4849

50+
cc_library(
51+
name = "grad_ops",
52+
deps = [
53+
":array_grad",
54+
":math_grad",
55+
],
56+
)
57+
58+
cc_library(
59+
name = "grad_testutil",
60+
testonly = 1,
61+
srcs = ["gradients/grad_testutil.cc"],
62+
hdrs = ["gradients/grad_testutil.h"],
63+
deps = [
64+
":grad_op_registry",
65+
":ops",
66+
":scope",
67+
],
68+
)
69+
4970
cc_library(
5071
name = "ops",
5172
srcs = ["framework/ops.cc"],
@@ -146,6 +167,33 @@ cc_library(
146167
],
147168
)
148169

170+
cc_library(
171+
name = "array_grad",
172+
srcs = ["gradients/array_grad.cc"],
173+
deps = [
174+
":cc_ops",
175+
":grad_op_registry",
176+
":ops",
177+
":scope",
178+
"//tensorflow/core:core_cpu",
179+
"//tensorflow/core:framework",
180+
],
181+
)
182+
183+
tf_cc_test(
184+
name = "gradients/array_grad_test",
185+
deps = [
186+
":array_grad",
187+
":cc_ops",
188+
":grad_op_registry",
189+
":grad_testutil",
190+
":testutil",
191+
"//tensorflow/core:test",
192+
"//tensorflow/core:test_main",
193+
"//tensorflow/core:testlib",
194+
],
195+
)
196+
149197
cc_library(
150198
name = "math_grad",
151199
srcs = ["gradients/math_grad.cc"],
@@ -164,12 +212,10 @@ tf_cc_test(
164212
deps = [
165213
":cc_ops",
166214
":grad_op_registry",
215+
":grad_testutil",
167216
":math_grad",
168-
"//tensorflow/core:all_kernels",
169-
"//tensorflow/core:core_cpu",
170-
"//tensorflow/core:framework",
217+
":testutil",
171218
"//tensorflow/core:lib_internal",
172-
"//tensorflow/core:tensorflow",
173219
"//tensorflow/core:test",
174220
"//tensorflow/core:test_main",
175221
"//tensorflow/core:testlib",
@@ -212,9 +258,7 @@ tf_cc_test(
212258
":client_session",
213259
":test_op",
214260
":test_op_op_lib",
215-
"//tensorflow/core:all_kernels",
216-
"//tensorflow/core:core_cpu",
217-
"//tensorflow/core:tensorflow",
261+
":testutil",
218262
"//tensorflow/core:test",
219263
"//tensorflow/core:test_main",
220264
"//tensorflow/core:testlib",
@@ -254,6 +298,23 @@ cc_library(
254298
alwayslink = 1,
255299
)
256300

301+
cc_library(
302+
name = "testutil",
303+
testonly = 1,
304+
srcs = ["framework/testutil.cc"],
305+
hdrs = ["framework/testutil.h"],
306+
deps = [
307+
":client_session",
308+
":ops",
309+
":scope",
310+
"//tensorflow/core:all_kernels",
311+
"//tensorflow/core:core_cpu",
312+
"//tensorflow/core:lib_internal",
313+
"//tensorflow/core:tensorflow",
314+
"//tensorflow/core:testlib",
315+
],
316+
)
317+
257318
tf_gen_op_wrappers_cc(
258319
name = "test_op",
259320
op_lib_names = [

‎tensorflow/cc/framework/cc_ops_test.cc‎

Lines changed: 15 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,11 @@ limitations under the License.
1414
==============================================================================*/
1515

1616
#include "tensorflow/cc/client/client_session.h"
17+
#include "tensorflow/cc/framework/testutil.h"
1718
#include "tensorflow/cc/ops/standard_ops.h"
1819
#include "tensorflow/cc/ops/test_op.h"
1920
#include "tensorflow/core/framework/tensor_testutil.h"
20-
#include "tensorflow/core/graph/default_device.h"
2121
#include "tensorflow/core/lib/core/status_test_util.h"
22-
#include "tensorflow/core/platform/test.h"
2322

2423
namespace tensorflow {
2524
using namespace ops; // NOLINT(build/namespaces)
@@ -32,18 +31,6 @@ Output Linear(const Scope& scope, Input x, Input w, Input b) {
3231
return BiasAdd(cop_scopes.last, m, b);
3332
}
3433

35-
void GetTensors(const Scope& scope, OutputList tensors,
36-
std::vector<Tensor>* out) {
37-
ClientSession session(scope);
38-
TF_CHECK_OK(session.Run(tensors, out));
39-
}
40-
41-
void GetTensor(const Scope& scope, Output tensor, Tensor* out) {
42-
std::vector<Tensor> outputs;
43-
GetTensors(scope, {tensor}, &outputs);
44-
*out = outputs[0];
45-
}
46-
4734
void GetColocationConstraints(Output tensor, std::vector<string>* constraints) {
4835
constraints->clear();
4936
const auto& attrs = tensor.op().node()->def().attr();
@@ -73,7 +60,7 @@ TEST(CCOpTest, Basic) {
7360
MatMul m(root, c, {{41}, {1}});
7461
TF_EXPECT_OK(root.status());
7562
Tensor out;
76-
GetTensor(root, m, &out);
63+
test::GetTensor(root, m, &out);
7764
test::ExpectTensorEqual<int>(out, test::AsTensor<int>({42}, {1, 1}));
7865
}
7966

@@ -82,7 +69,7 @@ TEST(CCOpTest, Attrs) {
8269
auto m = MatMul(root, {{1}, {1}}, {{41}, {1}}, MatMul::TransposeA(true));
8370
TF_EXPECT_OK(root.status());
8471
Tensor out;
85-
GetTensor(root, m, &out);
72+
test::GetTensor(root, m, &out);
8673
test::ExpectTensorEqual<int>(out, test::AsTensor<int>({42}, {1, 1}));
8774
}
8875

@@ -92,7 +79,7 @@ TEST(CCOpTest, SplitConcat) {
9279
auto c = Concat(root, 0, {p[0], p[1]});
9380
TF_EXPECT_OK(root.status());
9481
Tensor out;
95-
GetTensor(root, c, &out);
82+
test::GetTensor(root, c, &out);
9683
test::ExpectTensorEqual<int>(out, test::AsTensor<int>({1, 2}, {2, 1}));
9784
}
9885

@@ -103,15 +90,15 @@ TEST(CCOpTest, CompositeOp) {
10390
TF_EXPECT_OK(root.status());
10491
EXPECT_EQ(l.node()->name(), "layer0");
10592
Tensor out;
106-
GetTensor(root, l, &out);
93+
test::GetTensor(root, l, &out);
10794
test::ExpectClose(out, test::AsTensor<float>({-0.3, 34.2}, {1, 2}));
10895
}
10996

11097
TEST(CCOpTest, MultiOutput) {
11198
Scope root = Scope::NewRootScope();
11299
auto u = Unique(root, {1, 2, 2, 4, 3, 2});
113100
std::vector<Tensor> outputs;
114-
GetTensors(root, {u.y, u.idx}, &outputs);
101+
test::GetTensors(root, {u.y, u.idx}, &outputs);
115102
test::ExpectTensorEqual<int>(outputs[0], test::AsTensor<int>({1, 2, 4, 3}));
116103
test::ExpectTensorEqual<int>(outputs[1],
117104
test::AsTensor<int>({0, 1, 1, 2, 3, 1}));
@@ -134,7 +121,7 @@ TEST(CCOpTest, ExampleTrainer) {
134121
// y_normalized = y ./ y_norm
135122
auto y_normalized = Div(root.WithOpName("y_normalized"), y, y_norm);
136123
Tensor out;
137-
GetTensor(root, y_normalized, &out);
124+
test::GetTensor(root, y_normalized, &out);
138125
test::ExpectTensorNear<float>(
139126
out, test::AsTensor<float>({0.98058069, -0.19611613}, {2, 1}), 1e-5);
140127
}
@@ -161,15 +148,15 @@ TEST(CCOpTest, ControlDeps) {
161148

162149
std::vector<Tensor> out;
163150

164-
GetTensors(root, {add}, &out);
151+
test::GetTensors(root, {add}, &out);
165152
test::ExpectTensorNear<float>(out[0], test::AsTensor<float>({42.0f}, {}),
166153
1e-5);
167154

168155
out.clear();
169156
// Note : GetTensors creates a new session, so 'v' is uninitialized.
170157
// sub should have no control deps, so it should not cause the assign to run.
171158
// Hence is_inited should be false.
172-
GetTensors(root, {sub, is_inited}, &out);
159+
test::GetTensors(root, {sub, is_inited}, &out);
173160
test::ExpectTensorNear<float>(out[0], test::AsTensor<float>({1.0f}, {}),
174161
1e-5);
175162
test::ExpectTensorEqual<bool>(out[1], test::AsTensor<bool>({false}, {}));
@@ -220,12 +207,12 @@ TEST(CCOpTest, TemplatedConst) {
220207
TF_EXPECT_OK(root.status());
221208

222209
Tensor out;
223-
GetTensor(root, c1, &out);
210+
test::GetTensor(root, c1, &out);
224211
test::ExpectTensorEqual<float>(
225212
out, test::AsTensor<float>({3.f, 2.f, -1.f, 0.f}, {2, 2}));
226213

227214
auto c2 = ops::Const<string>(root, {{"this"}, {"is"}, {"a"}, {"constant"}});
228-
GetTensor(root, c2, &out);
215+
test::GetTensor(root, c2, &out);
229216
test::ExpectTensorEqual<string>(
230217
out, test::AsTensor<string>({"this", "is", "a", "constant"}, {4, 1}));
231218
}
@@ -237,22 +224,22 @@ TEST(CCOpTest, EmptyConst) {
237224
TF_CHECK_OK(root.status());
238225

239226
Tensor out;
240-
GetTensor(root, c1, &out);
227+
test::GetTensor(root, c1, &out);
241228
test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {0}));
242229

243230
auto c2 = ops::Const(root, {{}});
244231
TF_CHECK_OK(root.status());
245-
GetTensor(root, c2, &out);
232+
test::GetTensor(root, c2, &out);
246233
test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {1, 0}));
247234

248235
auto c3 = ops::Const(root, {{{}, {}}});
249236
TF_CHECK_OK(root.status());
250-
GetTensor(root, c3, &out);
237+
test::GetTensor(root, c3, &out);
251238
test::ExpectTensorEqual<float>(out, Tensor(DT_FLOAT, {1, 2, 0}));
252239

253240
auto c4 = ops::Const<int>(root, {{{}}});
254241
TF_CHECK_OK(root.status());
255-
GetTensor(root, c4, &out);
242+
test::GetTensor(root, c4, &out);
256243
test::ExpectTensorEqual<int>(out, Tensor(DT_INT32, {1, 1, 0}));
257244

258245
ops::Const(root, {{}, {{}}});

‎tensorflow/cc/framework/gradients.cc‎

Lines changed: 27 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -111,8 +111,8 @@ class SymbolicGradientBuilder {
111111
std::unordered_set<int> output_nodes_;
112112

113113
// The set of node ids in `inputs_`. Used to identify nodes at backprop
114-
// frontier.
115-
std::unordered_set<int> input_nodes_;
114+
// frontier. Maps from Output -> index into `grad_outputs_`.
115+
std::unordered_map<Output, int, OutputHash, OutputEq> input_nodes_;
116116

117117
TF_DISALLOW_COPY_AND_ASSIGN(SymbolicGradientBuilder);
118118
};
@@ -173,10 +173,10 @@ Status SymbolicGradientBuilder::Initialize() {
173173
for (int i = 0; i < outputs_.size(); ++i) {
174174
output_nodes_.insert(outputs_[i].node()->id());
175175
}
176-
// Populate `input_nodes_` from node ids in `inputs_`.
176+
// Populate `input_nodes_` from Outputs in `inputs_`.
177177
input_nodes_.reserve(inputs_.size());
178178
for (int i = 0; i < inputs_.size(); ++i) {
179-
input_nodes_.insert(inputs_[i].node()->id());
179+
input_nodes_.insert({inputs_[i], i});
180180
}
181181

182182
// TODO(andydavis) Consider a more efficient data structure for `pending_` to
@@ -187,8 +187,10 @@ Status SymbolicGradientBuilder::Initialize() {
187187
std::unordered_set<Node*> visited;
188188
std::deque<Node*> queue;
189189
for (const Output& nout : inputs_) {
190-
queue.push_back(nout.node());
191-
visited.insert(nout.node());
190+
if (visited.find(nout.node()) == visited.end()) {
191+
queue.push_back(nout.node());
192+
visited.insert(nout.node());
193+
}
192194
}
193195

194196
// Going forward to figure out which endpoints need backprop-ed.
@@ -286,19 +288,31 @@ Status SymbolicGradientBuilder::AddGradients() {
286288
Node* n = ready_.front();
287289
ready_.pop_front();
288290

289-
// Check if `n` is a member of `input_nodes_` where we terminate backprop.
290-
auto iter = input_nodes_.find(n->id());
291-
if (iter != input_nodes_.end()) {
292-
// Stop backprop.
293-
continue;
294-
}
295-
296291
// dy[i] is the sum of i-th output's backpropped gradients.
297292
const int num_y = n->num_outputs();
298293
dy.clear();
299294
dy.resize(num_y, {nullptr, 0});
300295
for (int i = 0; i < num_y; ++i) {
301296
TF_RETURN_IF_ERROR(SumGradients({n, i}, &dy[i]));
297+
auto iter = input_nodes_.find({n, i});
298+
if (iter != input_nodes_.end()) {
299+
// Return gradients for Output in 'grad_outputs_'.
300+
(*grad_outputs_)[iter->second] = dy[i];
301+
}
302+
}
303+
304+
// Stop backprop if none of the inputs to `n` are in `backprops_'.
305+
bool stop_node = true;
306+
for (const Edge* e : n->in_edges()) {
307+
if (e->IsControlEdge()) continue;
308+
if (backprops_.find({e->src(), e->src_output()}) != backprops_.end()) {
309+
stop_node = false;
310+
break;
311+
}
312+
}
313+
314+
if (stop_node) {
315+
continue;
302316
}
303317

304318
if (IsPrimitiveOpWithNoGrad(n->type_string())) {
@@ -331,12 +345,6 @@ Status SymbolicGradientBuilder::AddGradients() {
331345
BackpropAlongEdge(dx[dx_index++], {e->src(), e->src_output()}));
332346
}
333347
}
334-
335-
// Return gradients for `inputs_` in `grad_outputs_`.
336-
for (int i = 0; i < inputs_.size(); ++i) {
337-
TF_RETURN_IF_ERROR(SumGradients(inputs_[i], &(*grad_outputs_)[i]));
338-
}
339-
340348
return Status::OK();
341349
}
342350

0 commit comments

Comments
 (0)