本記事は、連載記事 TensorFlow内部構造解析 の1つで、Python APIでグラフを構築する処理の内部について、ソースコードレベルで説明した記事になります。
TensorFlowを使って機械学習/深層学習の演算を行う場合、TensorFlowが提供するAPI 1 を使って演算を定義します。演算を定義するために、ユーザは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 を足し算するグラフが作られて実行され 2、4.1 が出力されます。
上記のサンプルプログラムは非常に単純なものですが、実際にTensorFlowのPython APIで行っていることは複雑です。本記事では、Python APIの処理のうち、グラフを構築する処理について、TensorFlowの動作を解明していきます。
tf.constant
まずはじめに、tf.constant の関数定義を見てみましょう。tf.constant は、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 クラスのインスタンスを生成します。
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 は、ユーザがこれまでに定義した計算グラフに対して、ノードを作って追加する処理となります。
@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 の定義を示します。
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);
}
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() の関数定義を見てみると、以下のようになっています。
@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自動生成（未執筆）で説明します。