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 "absl/strings/str_cat.h"
22 #include "tensorflow/core/lib/strings/proto_serialization.h"
23 #include "tensorflow/core/platform/env.h"
24 #include "tensorflow/core/platform/mutex.h"
25
26 namespace tensorflow {
27
28 namespace {
29
30 struct NameCounts {
31 mutex counts_mutex;
32 std::unordered_map<string, int> counts;
33 };
34
MakeUniqueFilename(string name)35 string MakeUniqueFilename(string name) {
36 static NameCounts& instance = *new NameCounts;
37
38 // Remove illegal characters from `name`.
39 for (int i = 0; i < name.size(); ++i) {
40 char ch = name[i];
41 if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?') {
42 name[i] = '_';
43 }
44 }
45
46 int count;
47 {
48 mutex_lock lock(instance.counts_mutex);
49 count = instance.counts[name]++;
50 }
51
52 string filename = name;
53 if (count > 0) {
54 absl::StrAppend(&filename, "_", count);
55 }
56 absl::StrAppend(&filename, ".pbtxt");
57 return filename;
58 }
59
60 #if defined(TENSORFLOW_LITE_PROTOS)
WriteToFile(const string & filepath,const::tensorflow::protobuf::MessageLite & proto)61 Status WriteToFile(const string& filepath,
62 const ::tensorflow::protobuf::MessageLite& proto) {
63 string s;
64 if (!SerializeToStringDeterministic(proto, &s)) {
65 return errors::Internal("Failed to serialize proto to string.");
66 }
67 return WriteStringToFile(Env::Default(), filepath, s);
68 }
69 #else
WriteToFile(const string & filepath,const::tensorflow::protobuf::Message & proto)70 Status WriteToFile(const string& filepath,
71 const ::tensorflow::protobuf::Message& proto) {
72 return WriteTextProto(Env::Default(), filepath, proto);
73 }
74 #endif
75
76 template <class T>
WriteTextProtoToUniqueFile(Env * env,const string & name,const char * proto_type,T & proto,const string & dirname)77 string WriteTextProtoToUniqueFile(Env* env, const string& name,
78 const char* proto_type, T& proto,
79 const string& dirname) {
80 const char* dir = nullptr;
81 if (!dirname.empty()) {
82 dir = dirname.c_str();
83 } else {
84 dir = getenv("TF_DUMP_GRAPH_PREFIX");
85 }
86 if (!dir) {
87 LOG(WARNING)
88 << "Failed to dump " << name << " because dump location is not "
89 << " specified through either TF_DUMP_GRAPH_PREFIX environment "
90 << "variable or function argument.";
91 return "(TF_DUMP_GRAPH_PREFIX not specified)";
92 }
93 Status status = env->RecursivelyCreateDir(dir);
94 if (!status.ok()) {
95 LOG(WARNING) << "Failed to create " << dir << " for dumping " << proto_type
96 << ": " << status;
97 return "(unavailable)";
98 }
99 string filepath = absl::StrCat(dir, "/", MakeUniqueFilename(name));
100 status = WriteToFile(filepath, proto);
101 if (!status.ok()) {
102 LOG(WARNING) << "Failed to dump " << proto_type << " to file: " << filepath
103 << " : " << status;
104 return "(unavailable)";
105 }
106 LOG(INFO) << "Dumped " << proto_type << " to " << filepath;
107 return filepath;
108 }
109
110 } // anonymous namespace
111
DumpGraphDefToFile(const string & name,GraphDef const & graph_def,const string & dirname)112 string DumpGraphDefToFile(const string& name, GraphDef const& graph_def,
113 const string& dirname) {
114 return WriteTextProtoToUniqueFile(Env::Default(), name, "GraphDef", graph_def,
115 dirname);
116 }
117
DumpGraphToFile(const string & name,Graph const & graph,const FunctionLibraryDefinition * flib_def,const string & dirname)118 string DumpGraphToFile(const string& name, Graph const& graph,
119 const FunctionLibraryDefinition* flib_def,
120 const string& dirname) {
121 GraphDef graph_def;
122 graph.ToGraphDef(&graph_def);
123 if (flib_def) {
124 *graph_def.mutable_library() = flib_def->ToProto();
125 }
126 return DumpGraphDefToFile(name, graph_def, dirname);
127 }
128
DumpFunctionDefToFile(const string & name,FunctionDef const & fdef,const string & dirname)129 string DumpFunctionDefToFile(const string& name, FunctionDef const& fdef,
130 const string& dirname) {
131 return WriteTextProtoToUniqueFile(Env::Default(), name, "FunctionDef", fdef,
132 dirname);
133 }
134
135 } // namespace tensorflow
136