本記事は、連載記事 TensorFlow内部構造解析の1つで、Python APIでグラフを構築する処理の内部について、ソースコードレベルで説明した記事になります。

TensorFlowを使って機械学習/深層学習の演算を行う場合、TensorFlowが提供するAPI ¹ を使って演算を定義します。演算を定義するために、ユーザはTensorFlowのAPIを呼び出しますが、この時TensorFlow内部で有向非巡回グラフ（DAG）の計算グラフが作られます。
そしてユーザは、演算を定義したあとにセッション（実行環境）を作成し、作成した計算グラフを Session::run() メソッドの引数に渡すことで、計算グラフで定義した演算を実行することができます。

例として、簡単なTensorFlowのサンプルを示します。本サンプルは、定数値 1.5 と 2.6 をTensorFlowを使って足し算するプログラムです。

import tensorflow as tf

# 計算グラフ構築
a = tf.constant(1.5)
b = tf.constant(2.6)
output = a + b

# 計算グラフ実行
with tf.Session() as sess:
  print(sess.run(output))

図に示すように、このサンプルでは、TensorFlowの内部で定数値 1.5 と 2.6 を足し算するグラフが作られて実行され ²、4.1 が出力されます。

上記のサンプルプログラムは非常に単純なものですが、実際にTensorFlowのPython APIで行っていることは複雑です。本記事では、Python APIの処理のうち、グラフを構築する処理について、TensorFlowの動作を解明していきます。

tf.constant

まずはじめに、tf.constant の関数定義を見てみましょう。tf.constant は、tensorflow/python/framework/constant_op.py に定義されています。

tensorflow/python/framework/constant_op.py

# 1. Python APIに登録
@tf_export("constant")
def constant(value, dtype=None, shape=None, name="Const", verify_shape=False):
  ...
  ctx = context.context()
  # 2. Graph Mode/Eager Modeの切り分け
  if ctx.executing_eagerly():
    ...
    return
  g = ops.get_default_graph()
  # 3. Constant値の設定
  tensor_value = attr_value_pb2.AttrValue()
  tensor_value.tensor.CopyFrom(
      tensor_util.make_tensor_proto(
          value, dtype=dtype, shape=shape, verify_shape=verify_shape))
  dtype_value = attr_value_pb2.AttrValue(type=tensor_value.tensor.dtype)
  # 4. ノードの追加
  const_tensor = g.create_op(
      "Const", [], [dtype_value.type],
      attrs={"value": tensor_value,
             "dtype": dtype_value},
      name=name).outputs[0]
  return const_tensor

tf.exportデコレータ

tf.constant の関数定義でまず最初に目につくのが、デコレータ @tf_export("constant") です。
デコレータ @tf_export で修飾されている関数は、引数に指定された名前を関数名とするTensorFlowのPython APIとして、ユーザのプログラムから利用できるようになります。つまり、ユーザプログラムから tf.constant を呼び出すことで、tensorflow/python/framework/constant_op.py に定義された constant 関数が実行されます。

Graph ModeとEager Mode

tf.constant の処理を見ていくと、最初に実行モードによって処理を切り分けるコードが存在します。TensorFlowには実行モードが2つ存在し、それぞれGraph ModeとEager Modeと呼ばれており、これらの2つの実行モードを切り替える処理が、tf.constant の最初で行う処理となります。

ここで、TensorFlowの2つの実行モードについて簡単に説明します。

1つ目のGraph Modeは、深層学習フレームワークの世界において、Symbolic ProgrammingやDefine-And-Runと呼ばれているものです。Define-And-Runという用語が示すとおり、計算グラフを作った後に（必要に応じて計算グラフを最適化し）、計算グラフで定義された演算を実行します。
TensorFlowは、デフォルトでGraph Modeが適用されています。

2つ目はEager Modeです。こちらは、深層学習フレームワークの世界において、Imperative ProgrammingやDefine-By-Runと呼ばれています。計算グラフを作ってから実行するGraph Modeとは異なり、計算グラフを作らずに逐次演算を実行して結果を返します。
TensorFlowは、デフォルトでGraph Modeが適用されているため、Eager Modeを利用する場合は手動で設定する必要があります。TensorFlowでEager Modeを有効化するためには、tf.enable_eager_execution を実行します。

TensorFlowのPython APIでは、ctx.executing_eagerly がTrueの場合にEager Modeであると判定します。最初に示したサンプルでは、tf.enable_eager_execution を実行していません。このため、TensorFlowのデフォルトであるGraph Modeが適用されており、ctx.executing_eagerly は False となります。

Constant値の設定

tf.constant は、定数値を生成するConstノードを作ります。tf.constant は、後述する tf.add や tf.sub などのように、ノードを作る処理という点では同じ処理になりますが、Constノードによって生成される定数値の定義もノードの情報として持つ必要があります。TensorFlowでは、この定数値をProtocol Bufferを使って定義します。Protocol Bufferについては、こちらの記事を参照してください。

定数値は、Protocol Buffer TensorProto によって表現され、tensor_util.make_tensor_proto で作られます。
tensor_util.make_tensor_proto の内部では以下のように、tensorflow/core/framework/tensor.proto から生成された tensor_pb2 モジュールを使って、TensorProto クラスのインスタンスを生成します。

tensorflow/python/framework/tensor_util.py

  tensor_proto = tensor_pb2.TensorProto(
      dtype=numpy_dtype.as_datatype_enum,
      tensor_shape=tensor_shape.as_shape(shape).as_proto())

定数値は、Constノードの属性値（Attribute）として設定され、その属性値は、Protocol Buffer AttrValue によって表現されます。
そこで、tensorflow/core/framework/attr_value.proto から生成された attr_value_pb2 モジュールを使って、AttrValue クラスのインスタンスを生成し、先ほど作成した tensor_proto の値を設定します。

ノードの追加

ノードに設定するAttributeの作成が完了したため、実際にノードを作ります。

ノードは、tf.Graph の create_op メソッドを呼ぶことによって作ることができます。tf.Graph.create_op メソッドは、tensorflow/python/framework/ops.py に定義されています。

create_op メソッドを呼び出すと、本メソッドを呼び出したグラフにノードが追加されますが、ここでは、ops.get_default_graph によって取得したグラフにノードが追加されます。ops.get_default_graph は、TensorFlowのPython APIとして提供されている tf.get_default_graph と同一の定義であり、ユーザがこれまでに定義した計算グラフ tf.Graph を取得することができます。
つまり create_op は、ユーザがこれまでに定義した計算グラフに対して、ノードを作って追加する処理となります。

tensorflow/python/framework/ops.py

@tf_export("Graph")
class Graph(object):
  def create_op(...):
# ...
    # 1. NodeDefの作成
    node_def = _NodeDef(op_type, name, device=None, attrs=attrs)
# ...
    # 2. NodeDefからC++層にグラフを作成する
    ret = Operation(...)
# ...
    return ret

さて create_op メソッドについて、もう少し掘り下げて説明します。

create_op メソッドは最初に、Protocol Buffersでグラフのノードを表現する NodeDef を作成します。続いて Operation クラスを作りますが、そのコンストラクタ内で TF_NewOperation などのC APIを呼び出し、C++におけるグラフのノードを表現する Node クラスを作ります。

このことを理解するために、TF_NewOperation の先で呼び出される TF_NewOperationLocked と、TF_FinishOperation の先で呼び出される TF_FinishOperationLocked の定義を示します。

tensorflow/c/c_api.cc

static TF_OperationDescription* TF_NewOperationLocked(TF_Graph* graph,
                                                      const char* op_type,
                                                      const char* oper_name)
    EXCLUSIVE_LOCKS_REQUIRED(graph->mu) {
  return new TF_OperationDescription(graph, op_type, oper_name);
}

tensorflow/c/c_api.cc

static TF_Operation* TF_FinishOperationLocked(TF_OperationDescription* desc,
                                              TF_Status* status)
EXCLUSIVE_LOCKS_REQUIRED(desc->graph->mu) {
// ...
    status->status = desc->node_builder.Finalize(&desc->graph->graph, &ret);
// ...
  return ToOperation(ret);
}

最初に TF_NewOperation は、TF_OperationDescription を作ります。先ほど、TF_NewOperation は Node クラスを作ると書きましたが、TF_OperationDescription は Node クラスを構築するための NodeBuilder を持っています。
そして、C APIを呼び出してノードのAttributeを設定し、最終的に NodeBuilder::Finalize メソッドを呼び出すことによって、Node クラスを構築します。
最後に ToOperation 関数を呼び出して、作成した Node クラスを TF_Operation に変更していますが、これは単純に Node クラス1つを内包した構造体ですので、実質 Node クラスと同様と考えてよいです。

create_op メソッドの戻り値の output 変数は、Tensor オブジェクトです。これが最終的に tf.constant の戻り値となり、他のTensorFlowのAPIの入力として利用されます。

tf.add

続いて output = a + b ですが、これは output = tf.add(a, b) と読み替えることができるため、tf.add() の処理を確認すればよいことになります。ここで、tf.add() の関数定義を見てみると、以下のようになっています。

/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/gen_math_ops.py

@tf_export('add')
def add(x, y, name=None):
  _ctx = _context.context()
  if not _ctx.executing_eagerly():
    # Graph Mode時の処理
    _, _, _op = _op_def_lib._apply_op_helper(
        "Add", x=x, y=y, name=name)
    _result = _op.outputs[:]
    _inputs_flat = _op.inputs
    _attrs = ("T", _op.get_attr("T"))
    _execute.record_gradient(
      "Add", _inputs_flat, _attrs, _result, name)
    _result, = _result
    return _result
  else:
    # Eager Mode時の処理 ...

tf.export デコレータの役割や、ctx.executing_eagerly を使ってGraph ModeとEager Modeを切り替えている点は、tf.constant と同様です。
ここでは、tf.constant で説明していなかった、残りの処理について解析してみます。

最初に _apply_op_helper について説明します。_apply_op_helper は、tf.constant の内部で呼ばれていた create_op に渡す引数を設定して create_op を呼び出す便利関数です。_apply_op_helper は、以下のような引数を受け取ります。

引数名	説明
`op_type_name`	Operation名
`name`	ユーザ指定のノード名
`keyword`	Operationの入力

今回の場合は、Operation名を Add とし、Operationの入力として x と y を受け取ります。ユーザ指定のノード名は、tf.add() を呼び出したときに指定された引数 name をそのまま渡します。

上記の処理により、ops.get_default_graph によって取得したグラフにAddノードが追加されます。

Operationの自動生成

今回示した tf.add と同様、tf.sub などの演算の関数定義を調べてみると、_apply_op_helper に渡す引数が演算によって変わるのみで、大まかな処理の流れはほぼ同じであることが確認できます。

さて、これらのAPIが定義されているソースコードは、GitHub上（ビルド前のソースコード）には存在せず、これらのAPIの定義を示したソースコードがどこに存在するのか疑問を持つかもしれません。結論から言うと tf.add 等の演算のAPIのソースコードは、あるパターンに沿って決まる処理であるため、実装のコストを省く目的から、BazelでTensorFlowをビルドした時に自動生成されます。

APIの自動生成の詳細については、Operation自動生成（未執筆）で説明します。

TensorFlowは、Python、C++、Go、Java向けにAPIを提供しています。 ↩
実際は最適化され、定数値 4.1 だけのノードになりますが、詳細は本記事では説明しません。 ↩