1 | /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
2 | |
3 | Licensed under the Apache License, Version 2.0 (the "License"); |
4 | you may not use this file except in compliance with the License. |
5 | You may obtain a copy of the License at |
6 | |
7 | http://www.apache.org/licenses/LICENSE-2.0 |
8 | |
9 | Unless required by applicable law or agreed to in writing, software |
10 | distributed under the License is distributed on an "AS IS" BASIS, |
11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | See the License for the specific language governing permissions and |
13 | limitations under the License. |
14 | ==============================================================================*/ |
15 | |
16 | // Helper functions for dumping Graphs, GraphDefs, and FunctionDefs to files for |
17 | // debugging. |
18 | |
19 | #include "tensorflow/core/util/dump_graph.h" |
20 | |
21 | #include <memory> |
22 | #include <unordered_map> |
23 | |
24 | #include "absl/strings/match.h" |
25 | #include "absl/strings/str_cat.h" |
26 | #include "tensorflow/core/lib/strings/proto_serialization.h" |
27 | #include "tensorflow/core/platform/env.h" |
28 | #include "tensorflow/core/platform/file_system.h" |
29 | #include "tensorflow/core/platform/mutex.h" |
30 | #include "tensorflow/core/platform/path.h" |
31 | #include "tensorflow/core/platform/strcat.h" |
32 | |
33 | namespace tensorflow { |
34 | |
35 | namespace { |
36 | using strings::StrCat; |
37 | |
38 | struct NameCounts { |
39 | mutex counts_mutex; |
40 | std::unordered_map<string, int> counts; |
41 | }; |
42 | |
43 | string MakeUniqueFilename(string name, const string& suffix = ".pbtxt" ) { |
44 | static NameCounts& instance = *new NameCounts; |
45 | |
46 | // Remove illegal characters from `name`. |
47 | for (int i = 0; i < name.size(); ++i) { |
48 | char ch = name[i]; |
49 | if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' || |
50 | ch == '\\') { |
51 | name[i] = '_'; |
52 | } |
53 | } |
54 | |
55 | int count; |
56 | { |
57 | mutex_lock lock(instance.counts_mutex); |
58 | count = instance.counts[name]++; |
59 | } |
60 | |
61 | string filename = name; |
62 | if (count > 0) { |
63 | absl::StrAppend(&filename, "_" , count); |
64 | } |
65 | absl::StrAppend(&filename, suffix); |
66 | return filename; |
67 | } |
68 | |
69 | struct GraphDumperConfig { |
70 | mutex mu; |
71 | |
72 | // The dumper and suffix configured. |
73 | struct Config { |
74 | bool IsSet() const { return dumper != nullptr; } |
75 | std::function<Status(const Graph& graph, |
76 | const FunctionLibraryDefinition* flib_def, |
77 | WritableFile*)> |
78 | dumper = nullptr; |
79 | string suffix = ".pbtxt" ; |
80 | } config TF_GUARDED_BY(mu); |
81 | |
82 | // Returns whether a custom dumper is set. |
83 | bool IsSet() TF_LOCKS_EXCLUDED(mu) { |
84 | mutex_lock lock(mu); |
85 | return config.IsSet(); |
86 | } |
87 | }; |
88 | |
89 | GraphDumperConfig& GetGraphDumperConfig() { |
90 | static GraphDumperConfig config; |
91 | return config; |
92 | } |
93 | |
94 | // WritableFile that simply prints to stderr. |
95 | class StderrWritableFile : public WritableFile { |
96 | public: |
97 | StderrWritableFile() {} |
98 | |
99 | Status Append(StringPiece data) override { |
100 | fprintf(stderr, "%.*s" , static_cast<int>(data.size()), data.data()); |
101 | return OkStatus(); |
102 | } |
103 | |
104 | Status Close() override { return OkStatus(); } |
105 | |
106 | Status Flush() override { |
107 | fflush(stderr); |
108 | return OkStatus(); |
109 | } |
110 | |
111 | Status Name(StringPiece* result) const override { |
112 | *result = "stderr" ; |
113 | return OkStatus(); |
114 | } |
115 | |
116 | Status Sync() override { return OkStatus(); } |
117 | |
118 | Status Tell(int64_t* position) override { |
119 | return errors::Unimplemented("Stream not seekable" ); |
120 | } |
121 | }; |
122 | |
123 | Status CreateWritableFile(Env* env, const string& dirname, const string& name, |
124 | const string& suffix, string* filepath, |
125 | std::unique_ptr<WritableFile>* file) { |
126 | string dir; |
127 | if (!dirname.empty()) { |
128 | dir = dirname; |
129 | } else { |
130 | const char* prefix = getenv("TF_DUMP_GRAPH_PREFIX" ); |
131 | if (prefix != nullptr) dir = prefix; |
132 | } |
133 | if (dir.empty()) { |
134 | LOG(WARNING) |
135 | << "Failed to dump " << name << " because dump location is not " |
136 | << " specified through either TF_DUMP_GRAPH_PREFIX environment " |
137 | << "variable or function argument." ; |
138 | return errors::InvalidArgument("TF_DUMP_GRAPH_PREFIX not specified" ); |
139 | } |
140 | |
141 | if (absl::EqualsIgnoreCase(dir, "sponge" ) || |
142 | absl::EqualsIgnoreCase(dir, "test_undeclared_outputs_dir" )) { |
143 | if (!io::GetTestUndeclaredOutputsDir(&dir)) { |
144 | LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge, but " |
145 | "TEST_UNDECLARED_OUTPUT_DIRS is not set, dumping to log" ; |
146 | dir = "-" ; |
147 | } |
148 | } |
149 | |
150 | *filepath = "NULL" ; |
151 | if (dir == "-" ) { |
152 | *file = std::make_unique<StderrWritableFile>(); |
153 | *filepath = "(stderr)" ; |
154 | return OkStatus(); |
155 | } |
156 | |
157 | TF_RETURN_IF_ERROR(env->RecursivelyCreateDir(dir)); |
158 | *filepath = io::JoinPath(dir, MakeUniqueFilename(name, suffix)); |
159 | return env->NewWritableFile(*filepath, file); |
160 | } |
161 | |
162 | Status WriteTextProtoToUniqueFile(const tensorflow::protobuf::Message& proto, |
163 | WritableFile* file) { |
164 | string s; |
165 | if (!::tensorflow::protobuf::TextFormat::PrintToString(proto, &s)) { |
166 | return errors::FailedPrecondition("Unable to convert proto to text." ); |
167 | } |
168 | TF_RETURN_IF_ERROR(file->Append(s)); |
169 | StringPiece name; |
170 | TF_RETURN_IF_ERROR(file->Name(&name)); |
171 | VLOG(5) << name; |
172 | VLOG(5) << s; |
173 | return file->Close(); |
174 | } |
175 | |
176 | Status WriteTextProtoToUniqueFile( |
177 | const tensorflow::protobuf::MessageLite& proto, WritableFile* file) { |
178 | string s; |
179 | if (!SerializeToStringDeterministic(proto, &s)) { |
180 | return errors::Internal("Failed to serialize proto to string." ); |
181 | } |
182 | StringPiece name; |
183 | TF_RETURN_IF_ERROR(file->Name(&name)); |
184 | VLOG(5) << name; |
185 | VLOG(5) << s; |
186 | TF_RETURN_IF_ERROR(file->Append(s)); |
187 | return file->Close(); |
188 | } |
189 | |
190 | string DumpToFile(const string& name, const string& dirname, |
191 | const string& suffix, const string& type_name, |
192 | std::function<Status(WritableFile*)> dumper) { |
193 | string filepath; |
194 | std::unique_ptr<WritableFile> file; |
195 | Status status = CreateWritableFile(Env::Default(), dirname, name, suffix, |
196 | &filepath, &file); |
197 | if (!status.ok()) { |
198 | return StrCat("(failed to create writable file: " , status.ToString(), ")" ); |
199 | } |
200 | |
201 | status = dumper(file.get()); |
202 | if (!status.ok()) { |
203 | return StrCat("(failed to dump " , type_name, " to '" , filepath, |
204 | "': " , status.ToString(), ")" ); |
205 | } |
206 | LOG(INFO) << "Dumped " << type_name << " to " << filepath; |
207 | return filepath; |
208 | } |
209 | |
210 | } // anonymous namespace |
211 | |
212 | void SetGraphDumper( |
213 | std::function<Status(const Graph& graph, |
214 | const FunctionLibraryDefinition* flib_def, |
215 | WritableFile*)> |
216 | dumper, |
217 | string suffix) { |
218 | GraphDumperConfig& dumper_config = GetGraphDumperConfig(); |
219 | mutex_lock lock(dumper_config.mu); |
220 | dumper_config.config.dumper = dumper; |
221 | dumper_config.config.suffix = suffix; |
222 | } |
223 | |
224 | string DumpGraphDefToFile(const string& name, GraphDef const& graph_def, |
225 | const string& dirname) { |
226 | return DumpToFile(name, dirname, ".pbtxt" , "Graph" , [&](WritableFile* file) { |
227 | return WriteTextProtoToUniqueFile(graph_def, file); |
228 | }); |
229 | } |
230 | |
231 | string DumpCostGraphDefToFile(const string& name, CostGraphDef const& graph_def, |
232 | const string& dirname) { |
233 | return DumpToFile(name, dirname, ".pbtxt" , "Graph" , [&](WritableFile* file) { |
234 | return WriteTextProtoToUniqueFile(graph_def, file); |
235 | }); |
236 | } |
237 | |
238 | string DumpGraphToFile(const string& name, Graph const& graph, |
239 | const FunctionLibraryDefinition* flib_def, |
240 | const string& dirname) { |
241 | auto& dumper_config = GetGraphDumperConfig(); |
242 | if (dumper_config.IsSet()) { |
243 | GraphDumperConfig::Config config; |
244 | { |
245 | mutex_lock lock(dumper_config.mu); |
246 | config = dumper_config.config; |
247 | } |
248 | if (config.IsSet()) { |
249 | return DumpToFile(name, dirname, config.suffix, "Graph" , |
250 | [&](WritableFile* file) { |
251 | return config.dumper(graph, flib_def, file); |
252 | }); |
253 | } |
254 | } |
255 | |
256 | GraphDef graph_def; |
257 | graph.ToGraphDef(&graph_def); |
258 | if (flib_def) { |
259 | *graph_def.mutable_library() = flib_def->ToProto(); |
260 | } |
261 | return DumpGraphDefToFile(name, graph_def, dirname); |
262 | } |
263 | |
264 | string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef, |
265 | const string& dirname) { |
266 | return DumpToFile(name, dirname, ".pbtxt" , "FunctionDef" , |
267 | [&](WritableFile* file) { |
268 | return WriteTextProtoToUniqueFile(fdef, file); |
269 | }); |
270 | } |
271 | |
272 | } // namespace tensorflow |
273 | |