Skip to content

Commit 00526b6

Browse files
caisqtensorflower-gardener
caisq
authored and
tensorflower-gardener
committed
tfdbg: Watch output tensors without emitted edges
Currently, DebugNodeInserter is capable of only watching node output slots that have data edges (i.e., non-control edges) emitting from them. However, it would be useful to watch node output slots that do not have edges emitting from them, or output slots in nodes that have only control edges emitting from them. A common case is in the optimization ops in the backward paths of a training graph, where the train ops receives control edges from a number of nodes that each update the value of a Variable. Each of these Variable-updating nodes has a data output slot (slot 0), but those slots are not connected to any data edges. In the watch list, if you specify variable_updating_op:0, this tensor will not get watched without this CL. This CL adds the ability to watch such nodes' outputs. The CL also makes it possible to watch output slots that do not have any control or non-control edges emitting from them. Change: 136825445
1 parent e9e56f2 commit 00526b6

3 files changed

Lines changed: 316 additions & 88 deletions

File tree

‎tensorflow/core/debug/debug_gateway_test.cc‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,108 @@ TEST_F(SessionDebugMinusAXTest,
448448
}
449449
}
450450

451+
class SessionDebugOutputSlotWithoutOngoingEdgeTest : public ::testing::Test {
452+
public:
453+
void Initialize() {
454+
Graph graph(OpRegistry::Global());
455+
456+
#if GOOGLE_CUDA
457+
const string kDeviceName = "/job:localhost/replica:0/task:0/gpu:0";
458+
#else
459+
const string kDeviceName = "/job:localhost/replica:0/task:0/cpu:0";
460+
#endif
461+
462+
Tensor a_tensor(DT_FLOAT, TensorShape({1, 1}));
463+
test::FillValues<float>(&a_tensor, {42.0});
464+
Node* a = test::graph::Constant(&graph, a_tensor);
465+
a->set_assigned_device_name(kDeviceName);
466+
467+
Node* c = test::graph::Constant(&graph, a_tensor);
468+
c->set_assigned_device_name(kDeviceName);
469+
c_ = c->name();
470+
471+
// Node c will be executed only because of the control edge from c to y.
472+
// Its output slot (slot 0) does not have an outgoing edge. This test
473+
// is for testing that the debugger can watch that slot properly.
474+
Node* y = test::graph::NoOp(&graph, {c});
475+
y->set_assigned_device_name(kDeviceName);
476+
y_ = y->name();
477+
478+
test::graph::ToGraphDef(&graph, &def_);
479+
}
480+
481+
string c_;
482+
string y_;
483+
GraphDef def_;
484+
};
485+
486+
TEST_F(SessionDebugOutputSlotWithoutOngoingEdgeTest,
487+
WatchSlotWithoutOutgoingEdge) {
488+
Initialize();
489+
std::unique_ptr<DirectSession> session(CreateSession());
490+
ASSERT_TRUE(session != nullptr);
491+
492+
DebugGateway debug_gateway(session.get());
493+
494+
// Supply completion and value callbacks
495+
mutex mu;
496+
497+
string debug_identity_node_name = DebugNodeInserter::GetDebugNodeName(
498+
strings::StrCat(c_, ":", 0), 0, "DebugIdentity");
499+
500+
Notification callbacks_done;
501+
502+
debug_gateway.SetNodeCompletionCallback(
503+
[&mu, &callbacks_done](const string& node_name, const bool any_output) {
504+
mutex_lock l(mu);
505+
if (node_name == "_SINK" && !callbacks_done.HasBeenNotified()) {
506+
callbacks_done.Notify();
507+
}
508+
});
509+
510+
std::vector<Tensor> debug_identity_tensor_vals;
511+
debug_gateway.SetNodeValueCallback(
512+
[this, &mu, &debug_identity_node_name, &debug_identity_tensor_vals](
513+
const string& node_name, const int output_slot,
514+
const Tensor& tensor_value, const bool is_ref) {
515+
mutex_lock l(mu);
516+
517+
if (node_name == debug_identity_node_name && output_slot == 0) {
518+
debug_identity_tensor_vals.push_back(tensor_value);
519+
}
520+
});
521+
522+
// Add DebugIdentity watch on c:0, which does not have an outgoing edge.
523+
RunOptions run_opts;
524+
run_opts.set_output_partition_graphs(true);
525+
526+
DebugTensorWatch* tensor_watch_opts = run_opts.add_debug_tensor_watch_opts();
527+
tensor_watch_opts->set_node_name(c_);
528+
tensor_watch_opts->set_output_slot(0);
529+
tensor_watch_opts->add_debug_ops("DebugIdentity");
530+
531+
TF_ASSERT_OK(session->Create(def_));
532+
533+
// Invoke Session::Run() on y.
534+
std::vector<std::pair<string, Tensor>> inputs;
535+
std::vector<string> output_names;
536+
std::vector<string> target_nodes = {y_};
537+
std::vector<Tensor> outputs;
538+
539+
RunMetadata run_metadata;
540+
Status s = session->Run(run_opts, inputs, output_names, target_nodes,
541+
&outputs, &run_metadata);
542+
TF_ASSERT_OK(s);
543+
544+
// Wait for callbacks to complete.
545+
callbacks_done.WaitForNotification();
546+
547+
// Assert that DebugIdentity node watching the control edge has been run.
548+
ASSERT_EQ(1, debug_identity_tensor_vals.size());
549+
auto mat_identity = debug_identity_tensor_vals[0].matrix<float>();
550+
ASSERT_EQ(42.0, mat_identity(0, 0));
551+
}
552+
451553
class SessionDebugVariableTest : public ::testing::Test {
452554
public:
453555
void Initialize() {

‎tensorflow/core/debug/debug_graph_utils.cc‎

Lines changed: 60 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -73,87 +73,69 @@ Status DebugNodeInserter::InsertNodes(
7373
}
7474

7575
DeviceType device_type = DeviceType{device->device_type()};
76-
// 1. Record existing edges in the graph.
77-
std::vector<const Edge*> existing_edges;
78-
for (const Edge* edge : graph->edges()) {
79-
existing_edges.push_back(edge);
80-
}
81-
82-
// A map from tensor names to edges to be removed
83-
std::unordered_map<string, std::vector<const Edge*>> edges_to_remove;
84-
// A map from tensor names to newly added debug nodes (maybe more than one
85-
// for a given tensor).
86-
std::unordered_map<string, std::vector<Node*>> added_debug_nodes;
87-
std::unordered_map<string, Node*> added_copy_nodes;
8876

89-
// 2. Iterate through the edges, look for edges that match the tensor watch
90-
// list.
91-
for (const Edge* edge : existing_edges) {
92-
Node* src_node = edge->src();
93-
Node* dst_node = edge->dst();
94-
95-
if (edge->IsControlEdge()) {
96-
continue;
77+
// Keep track of all edges to be removed.
78+
std::vector<const Edge*> edges_to_remove;
79+
80+
for (Node* src_node : graph->nodes()) {
81+
// Make a map from output slot to outgoing edges from the slot.
82+
std::unordered_map<int, std::vector<const Edge*>> output_slot_to_edges;
83+
for (const Edge* edge : src_node->out_edges()) {
84+
const int src_output = edge->src_output();
85+
if (output_slot_to_edges.find(src_output) == output_slot_to_edges.end()) {
86+
output_slot_to_edges[src_output] = {edge};
87+
} else {
88+
output_slot_to_edges[src_output].push_back(edge);
89+
}
9790
}
9891

99-
const bool is_ref = IsRefType(dst_node->input_type(edge->dst_input()));
100-
MemoryType memory_type;
101-
MemoryTypeForOutput(device_type, graph, src_node, edge->src_output(),
102-
&memory_type);
103-
104-
const string tensor_name =
105-
strings::StrCat(src_node->name(), ":", edge->src_output());
106-
if (tensor_watches.find(tensor_name) == tensor_watches.end()) {
107-
// Add debug nodes only for edges with matching source node and source
108-
// output slot.
109-
continue;
110-
}
92+
// Iterate through all output slots of the node.
93+
for (int src_output_slot = 0; src_output_slot < src_node->num_outputs();
94+
++src_output_slot) {
95+
const string tensor_name =
96+
strings::StrCat(src_node->name(), ":", src_output_slot);
97+
if (tensor_watches.find(tensor_name) == tensor_watches.end()) {
98+
// Add debug nodes only for edges with matching source node and source
99+
// output slot.
100+
continue;
101+
}
111102

112-
if (added_copy_nodes.find(tensor_name) == added_copy_nodes.end()) {
113-
// It is the first time an edge with this source tensor is encountered:
114-
// we will:
115-
// 1) Mark this edge as to be removed, iff the destination node has
116-
// non-Ref input
117-
// 2) Create a Copy node
103+
// Now we have encountered a watched tensor. We will:
104+
// 1) Mark this edge as to be removed, iff this is a non-Reference
105+
// tensor
106+
// 2) Create a Copy node for the tensor
118107
// 3) Add a new edge, from the source tensor to the Copy node
119108
// 4) Add a new edge, from the Copy node to the destination node, iff
120-
// the destination node has non-Ref input
109+
// this is a non-Reference tensor.
121110
// 5) Create all the requested debug nodes and their edges to the Copy
122111
// node.
123-
if (!is_ref) {
124-
std::vector<const Edge*> node_edges_to_remove;
125-
node_edges_to_remove.push_back(edge);
126-
edges_to_remove[tensor_name] = node_edges_to_remove;
127-
}
112+
// 6) Add control edges from the debug nodes to the destination nodes
113+
// to ensure that the tensors values exported by the debug nodes
114+
// to the debug URLs reflect the values before the execution of
115+
// the destination nodes.
128116

129-
const DataType src_dt = src_node->output_type(edge->src_output());
117+
const DataType src_dt = src_node->output_type(src_output_slot);
118+
MemoryType memory_type;
119+
MemoryTypeForOutput(device_type, graph, src_node, src_output_slot,
120+
&memory_type);
130121

131-
// Create the copy node.
122+
// Create the copy node for the watched tensor.
132123
Node* copy_node;
133124
Status copy_s = CreateCopyNode(
134125
graph, device_type, memory_type == HOST_MEMORY, src_node->name(),
135-
edge->src_output(), src_dt, tensor_name, &copy_node);
126+
src_output_slot, src_dt, tensor_name, &copy_node);
136127
if (!copy_s.ok()) {
137128
return Status(
138129
error::FAILED_PRECONDITION,
139130
strings::StrCat("Failed to create Copy/CopyHost node for tensor ",
140131
tensor_name, ", due to: ", copy_s.error_message()));
141132
}
142133

143-
// Record the added copy node for later use.
144-
added_copy_nodes[tensor_name] = copy_node;
145-
146134
// Add edge from watched tensor to the copy node.
147-
graph->AddEdge(src_node, edge->src_output(), copy_node, 0);
148-
149-
// Add edge from the copy node to the destination node, iff the
150-
// destination node has non-Ref input.
151-
if (!is_ref) {
152-
graph->AddEdge(copy_node, 0, dst_node, edge->dst_input());
153-
}
135+
graph->AddEdge(src_node, src_output_slot, copy_node, 0);
154136

155137
// Create all requested debug nodes and their edges to the Copy node.
156-
std::vector<Node*> node_added_debug_nodes;
138+
std::vector<Node*> debug_nodes;
157139
for (size_t i = 0; i < tensor_watches[tensor_name].size(); ++i) {
158140
const string& debug_op_name = tensor_watches[tensor_name][i];
159141

@@ -169,47 +151,37 @@ Status DebugNodeInserter::InsertNodes(
169151
debug_s.error_message()));
170152
}
171153

172-
node_added_debug_nodes.push_back(debug_node);
173-
174154
// Create edges from the Copy node to the debug node.
175155
graph->AddEdge(copy_node, 0, debug_node, 0);
176156

157+
debug_nodes.push_back(debug_node);
158+
}
159+
160+
// Is the output a reference?
161+
const bool is_ref = IsRefType(src_node->output_type(src_output_slot));
162+
163+
// Iterate through all outgoing edges attached to the slot.
164+
for (const Edge* edge : output_slot_to_edges[src_output_slot]) {
165+
// Mark the edge for removal.
166+
if (!is_ref) {
167+
edges_to_remove.push_back(edge);
168+
graph->AddEdge(copy_node, 0, edge->dst(), edge->dst_input());
169+
}
170+
177171
// Add control edges from the debug nodes to the destination node
178172
// to ensure that the debug nodes are executed before the destination
179173
// node.
180-
graph->AddEdge(debug_node, Graph::kControlSlot, dst_node,
181-
Graph::kControlSlot);
182-
}
183-
added_debug_nodes[tensor_name] = node_added_debug_nodes;
184-
} else {
185-
// It is not the first time an edge with this source is encountered.
186-
// We will do the following iff the destination node has non-Ref input
187-
// 1) Mark the edge for removal
188-
// 2) Create an edge from the copy node to the destination node
189-
// Iff the destination has Ref-input, the edge will not change.
190-
// Regardless of whether the destination has Ref-inpt, we will
191-
// 3) Add control edges from the already-created debug node(s) for the
192-
// watched tensor to the destination node.
193-
if (!is_ref) {
194-
edges_to_remove[tensor_name].push_back(edge);
195-
graph->AddEdge(added_copy_nodes[tensor_name], 0, dst_node,
196-
edge->dst_input());
197-
}
198-
199-
for (Node* debug_node : added_debug_nodes[tensor_name]) {
200-
graph->AddEdge(debug_node, Graph::kControlSlot, dst_node,
201-
Graph::kControlSlot);
174+
for (Node* debug_node : debug_nodes) {
175+
graph->AddEdge(debug_node, Graph::kControlSlot, edge->dst(),
176+
Graph::kControlSlot);
177+
}
202178
}
203179
}
204180
}
205181

206182
// Remove all edges marked for removal.
207-
for (auto it : edges_to_remove) {
208-
std::vector<const Edge*> edges = it.second;
209-
210-
for (const Edge* edge : edges) {
211-
graph->RemoveEdge(edge);
212-
}
183+
for (const Edge* edge : edges_to_remove) {
184+
graph->RemoveEdge(edge);
213185
}
214186

215187
return Status::OK();

0 commit comments

Comments
 (0)