tensorflow の GraphDef を JSON 形式でシリアライズ/デシリアライズする

背景

protobuf のバイナリ/テキスト形式を使いたくない.
onnx も protobuf 形式でつらい.

ちょっとしたグラフをいじるプログラムを書きたいとしても...

• C++ では protobuf リンクが面倒である.
• python でも protobuf モジュールインストールが面倒である.
• python で import tensorflow が遅くてつらい.

JSON で扱うことで, たとえばグラフの解析やテスト用データ生成を, C++ json.hpp https://github.com/nlohmann/json でヘッダーオンリーライブラリで処理したり, python 標準の json モジュールだけで処理したい.

JSON 形式でダンプする.

tensorflow の GraphDef オブジェクト(Protobuf オブジェクト?)であれば, MessageToJson で Protobuf の JSON 形式にシリアライズできます.

from google.protobuf.json_format import MessageToJson

from tensorflow.core.framework import graph_pb2
graph_def = graph_pb2.GraphDef()

...

# to JSON
j = MessageToJson(graph_def)

with open("output.json", "w") as o:
    o.write(j)

テンソル値などのバイナリデータは base64 でエンコードされます.
たとえば Const は以下のような感じの JSON になります.

{ "node": [
    {
      "attr": {
        "dtype": {
          "type": "DT_FLOAT"
        },
        "value": {
          "tensor": {
            "dtype": "DT_FLOAT",
        "tensorContent": "AACAPwAAAEDNzAxAAACAQA==",
        "tensorShape": {
          "dim": [
            {
              "size": "4"
            }
          ]
        }
          }
        }
      },
      "name": "const1",
      "op": "Const"
    }
    ]
}

JSON を処理する

普通の JSON なので, いろいろ簡単に料理できます!(レイヤーの名前を変えたりとか, 解析とか)

Tensorflow で JSON データを読み戻す.

JsonToMessage は何故か無い.
JsonStringToMessage は C++ レイヤーではあるようですが, python binding では無いようです...

google.protobuf.json_format.Parse を使います.

import tensorflow as tf
from google.protobuf import json_format
from tensorflow.core.framework import graph_pb2

with open(sys.argv[1], "r") as f:
    graph_def = json_format.Parse(f.read(), graph_pb2.GraphDef())

    with tf.Session() as sess:
        tf.import_graph_def(graph_def)

Voila! めでたく JSON から復元できました.

TODO

• 優秀な Protobuf -> JSON 若人が, 人類史上最速で育成されるスキームを確立する旅に出たい.
• glTF のように, JSON で定義された業界標準の機械学習フォーマット本当にはよ!