Function InsertLogging

tensorflow/tools/graph_transforms/insert_logging.cc:31–146

Clears the device field of all ops in the graph.

// Clears the device field of all ops in the graph.
Status InsertLogging(const GraphDef& input_graph_def,
                     const TransformFuncContext& context,
                     GraphDef* output_graph_def) {
  std::unordered_set<string> ops;
  bool has_ops;
  if (context.params.count("op")) {
    has_ops = true;
    for (const string& op : context.params.at("op")) {
      ops.insert(op);
    }
  } else {
    has_ops = false;
  }

  std::unordered_set<string> prefixes;
  bool has_prefixes;
  if (context.params.count("prefix")) {
    has_prefixes = true;
    for (const string& prefix : context.params.at("prefix")) {
      prefixes.insert(prefix);
    }
  } else {
    has_prefixes = false;
  }

  string message;
  TF_RETURN_IF_ERROR(context.GetOneStringParameter("message", "", &message));

  bool show_name;
  TF_RETURN_IF_ERROR(
      context.GetOneBoolParameter("show_name", false, &show_name));

  bool show_op;
  TF_RETURN_IF_ERROR(context.GetOneBoolParameter("show_op", false, &show_op));

  int32 first_n;
  TF_RETURN_IF_ERROR(context.GetOneInt32Parameter("first_n", -1, &first_n));

  int32 summarize;
  TF_RETURN_IF_ERROR(
      context.GetOneInt32Parameter("summarize", 1024, &summarize));

  std::unordered_map<string, std::set<int>> node_outputs;
  for (const NodeDef& node : input_graph_def.node()) {
    for (const string& input : node.input()) {
      const string canonical_input = CanonicalInputName(input);
      string prefix;
      string name;
      string suffix;
      NodeNamePartsFromInput(canonical_input, &prefix, &name, &suffix);
      const string output_index_string = suffix.substr(1, suffix.size() - 1);
      int32 output_index;
      if (!strings::safe_strto32(output_index_string, &output_index)) {
        return errors::InvalidArgument("Couldn't understand output number in ",
                                       input);
      }
      node_outputs[name].insert(output_index);
    }