TensorFlow内部構造解析 (2.1) Python層におけるグラフ構築

本記事は、連載記事 [TensorFlow内部構造解析](https://qiita.com/AtuNuka/items/028780aab338d9b1a7c2) の1つで、Python APIでグラフを構築する処理の内部について、ソースコードレベルで説明した記事になります。

___

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

[^1]: TensorFlowは、Python、C++、Go、Java向けにAPIを提供しています。

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

```python
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``` を足し算するグラフが作られて実行され [^2]、```4.1``` が出力されます。

[^2]: 実際は最適化され、定数値 ```4.1``` だけのノードになりますが、詳細は本記事では説明しません。

![cg.png](https://qiita-image-store.s3.amazonaws.com/0/229755/8d43aebb-2e1f-7ad9-ca98-6abe15b14a8f.png)


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

# tf.constant

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

```python: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については、[こちらの記事](https://qiita.com/AtuNuka/items/8b6a9d632641fbc787d4) を参照してください。

定数値は、Protocol Buffer [TensorProto](https://qiita.com/AtuNuka/items/8b6a9d632641fbc787d4#tensorproto) によって表現され、```tensor_util.make_tensor_proto``` で作られます。
```tensor_util.make_tensor_proto``` の内部では以下のように、```tensorflow/core/framework/tensor.proto``` から生成された ```tensor_pb2``` モジュールを使って、```TensorProto``` クラスのインスタンスを生成します。

```python: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](https://qiita.com/AtuNuka/items/8b6a9d632641fbc787d4#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``` は、ユーザがこれまでに定義した計算グラフに対して、ノードを作って追加する処理となります。


```python: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``` の定義を示します。

```c: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);
}
```

```c: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``` クラスと同様と考えてよいです。

![python-c-graph.png](https://qiita-image-store.s3.amazonaws.com/0/229755/79d871d3-2a3a-618a-03b0-e54104561e77.png)

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


# tf.add

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

```python:/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自動生成（未執筆）で説明します。