Getting Protobuf Message for every TensorFlow Operation

Question

From here,

There are a few ways to get a list of the OpDefs for the registered ops:

TF_GetAllOpList in the C API retrieves all registered OpDef protocol messages. This can be used to write the generator in the client language. This requires that the client language have protocol buffer support in order to interpret the OpDef messages.

The C++ function OpRegistry::Global()->GetRegisteredOps() returns the same list of all registered OpDefs (defined in [tensorflow/core/framework/op.h]). This can be used to write the generator in C++ (particularly useful for languages that do not have protocol buffer support).

The ASCII-serialized version of that list is periodically checked in to [tensorflow/core/ops/ops.pbtxt] by an automated process.

But alas, I want to do this in Python like,

import tensorflow as tf
from google.protobuf import json_format
json_string = json_format.MessageToJson(tf.GetAllOpsList())

I want a way to get the Protobuf message for every operation in Tensorflow so that I can dump it as JSON via

score 2 · Accepted Answer · edited May 23 '17 at 10:29

It's in ops.txt. Here's an example of listing all OpDef messages for ops that produce strings outputs.

import tensorflow as tf

from tensorflow.core.framework import op_def_pb2
from google.protobuf import text_format

def get_op_types(op):
    for attr in op.attr:
        if attr.type != 'type':
            continue
        return list(attr.allowed_values.list.type)
    return []

# directory where you did "git clone"
tensorflow_git_base = "/Users/yaroslav/tensorflow.git"
ops_file = tensorflow_git_base+"/tensorflow/tensorflow/core/ops/ops.pbtxt"
ops = op_def_pb2.OpList()
text_format.Merge(open(ops_file).read(), ops)

for op in ops.op:
    # get templated string types
    if tf.string in get_op_types(op):
        print(op.name, op.summary)
    #for arg in op.input_arg:
    for arg in op.output_arg:
        if arg.type == tf.string:
            print(op.name, op.summary)
            break

** Added ** If you want to be sensitive to new ops being added you could reverse engineer how current Python wrappers do it. For instance, consider gen_array_ops.py file. It has the following snippet

def _InitOpDefLibrary():
  op_list = _op_def_pb2.OpList()
  _text_format.Merge(_InitOpDefLibrary.op_list_ascii, op_list)
  _op_def_registry.register_op_list(op_list)
  op_def_lib = _op_def_library.OpDefLibrary()
  op_def_lib.add_op_list(op_list)
  return op_def_lib


_InitOpDefLibrary.op_list_ascii = """op {
  name: "BatchMatrixBandPart"
  input_arg {
    name: "input"
    type_attr: "T"
  }
  input_arg {
    name: "num_lower"
    type: DT_INT64
  }
  input_arg {
    name: "num_upper"
    type: DT_INT64
  }
  output_arg {
    name: "band"
    type_attr: "T"
  }
  attr {
    name: "T"
    type: "type"
  }
  deprecation {
    version: 14
    explanation: "Use MatrixBandPart"
  }
}

So those messages protobufs are generated from underlying C code during generation of gen_array_ops. For tracking down how they were generated, see https://stackoverflow.com/a/41149557/419116

I would prefer a Python way so I can add custom operations and get the Protobuf out from those — , Feb 10 '17 at 17:27
Added extra info. You could either regenerate "ops.txt" by calling the same generation script as done during gen_*_ops.py generation, or find underlying C functionality and add some swig wrappers — Yaroslav Bulatov, Feb 10 '17 at 18:40
Thanks! Almost there. I don't care about ascii, but I need JSON. So I changed it to `from google.protobuf import json_format` and discovered that `json_format.MessageToJson(gen_array_ops._InitOpDefLibrary()._ops['Const'].op_def)` almost gives me what I want. This is just the metaInfoDef, but I need the graphDef block too — , Feb 10 '17 at 19:25
I got a bit further. Still stuck. http://stackoverflow.com/questions/42168977/tensorflow-manual-construction-of-graphdef — , Feb 10 '17 at 21:31

Getting Protobuf Message for every TensorFlow Operation

1 Answers1