計算グラフは数学的には有向非巡回グラフ（DAG）として表現され、ニューラルネットワークモデルを一貫した形式で抽象化できます。しかし、プログラミングでは条件分岐やループなどの制御フロー構造が頻繁に出現します。このような制御フローを計算グラフ上でどう表現するかは、計算グラフの重要な課題です。制御フローの導入により、開発者は分岐選択やループ制御を計算グラフに組み込むことができ、より複雑なニューラルネットワーク構造を構築できるようになります。
現在の多くのフレームワークでは、例えばPyTorchはPythonレベルでの制御フローのみをサポートしており、計算グラフ内部では制御フローは扱われません。このため、制御フローを含むモデルをデプロイする際には困難が生じます。特にPythonをサポートしない環境への展開が難しい問題があります。
計算グラフにおける制御フローと、コンパイラ理論でいう制御フローグラフ（CFG）は異なる概念です。CFGはプログラムの実行中に考えられるすべての実行パスを表す図であり、Frances E. Allenによって発見されました。CFGはコンパイラの最適化や静的解析に欠かせない要素です。
制御フローの背景
コンピュータサイエンスにおいて、制御フローとは命令や関数呼び出しなどの実行順序を定義するものです。たとえば、ある関数Aの出力に基づいて、関数BまたはCのどちらかを実行するといったケースです。
AIフレームワークは、ニューラルネットワークの計算を有向非巡回グラフ（DAG）として抽象化します。これはアルゴリズム開発者がニューラルネットワークを理解する方法と一致しており、学習特性に影響を与える演算子間のトポロジー構造を簡単に記述できます。しかし、近年のニューラルネットワーク構造の急速な進化により、純粋な計算グラフとして表現しきれない新しい構造が増加しています。
Transformer構造のニューラルネットワークを例に挙げましょう。Transformerはアテンションメカニズムに基づく構造で、複数のエンコーダーとデコーダーを積み重ねて構成されます。このような構造を自然に表現するには、ループ制御が必要です。Transformerのアルゴリズムを記述する際、`for`ループが頻繁に現れます。
課題
制御フローの導入により、計算グラフの構築と前方伝播に大きな変化が生じます。
- 計算グラフが動的に変化し、分岐やループの実行は実行時にデータ入力に基づいて決定されるため、静的な構造ではなく動的な構造となります。
- ループ制御フローの実装も困難です。ループが導入されることで、計算グラフに閉路が生じるため、トポロジカルソートができなくなります。そのため、前方向計算と逆方向計算の実装では、トポロジカルソートを放棄するか、ループを展開するなどの代替手段が必要です。
制御フローを持つニューラルネットワークモデルをサポートするために、AIフレームワークは動的な制御フロー構造をサポートする必要があります。現在の計算グラフベースの解決策には、以下の3つのアプローチがあります：
- ホスト言語の再利用：フロントエンドの制御フロー構造を再利用し、Pythonなどの制御ロジックでバックエンドの計算グラフを実行する。
- 制御フロー原語のサポート：AIフレームワークのバックエンドが制御フロー構造を直接サポートし、計算グラフ内で制御フローと計算フローを混在させる。
- ソースコード解析：フロントエンドで高級言語のコードを解析して計算グラフに変換し、バックエンドで制御フロー構造をサブグラフとして展開する。
動的グラフ
ニューラルネットワークモデルを実行するたびに、AIフレームワークはフロントエンドのプログラミング言語に基づいて一時的な計算グラフ（単一演算子）を動的に生成します。この方式では、計算グラフは動的に生成され、柔軟性があり、ネットワーク構造の変更に迅速に対応できます。この方式は動的計算グラフと呼ばれます。
ホスト言語の再利用
PyTorchは動的グラフを使用しており、これによりデバッグが容易になります。各ステップを詳細に制御・確認でき、イテレーションごとにネットワーク全体を再構築することも可能です。
PyTorchでは、Pythonの制御フロー構造をそのまま利用して動的グラフを構築します。制御フローが現れた場合、Pythonで制御が実行され、PyTorchのAPIが呼び出された段階で動的グラフが実行されます。
以下はTransformerのデコーダーにおける動的グラフの例です。forループで複数のデコーダーモジュールをスタックし、if文で正則化層を追加する構造です。
from torch import nn
class TransformerDecoder(nn.Module):
def __init__(self, decoder_layer, num_layers, norm=None):
super().__init__()
self.layers = _get_clones(decoder_layer, num_layers)
self.num_layers = num_layers
self.norm = norm
def forward(self, tgt: Tensor,
memory: Tensor,
tgt_mask: Optional[Tensor] = None,
memory_mask: Optional[Tensor] = None,
tgt_key_padding_mask: Optional[Tensor] = None,
memory_key_padding_mask: Optional[Tensor] = None,
tgt_is_causal: Optional[bool] = None,
memory_is_causal: bool = False):
seq_len = _get_seq_len(tgt, self.layers[0].self_attn.batch_first)
tgt_is_causal = _detect_is_causal_mask(tgt_mask, tgt_is_causal, seq_len)
for mod in self.layers:
output = mod(output, memory, tgt_mask=tgt_mask,
memory_mask=memory_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=memory_key_padding_mask,
tgt_is_causal=tgt_is_causal,
memory_is_causal=memory_is_causal)
if self.norm is not None:
output = self.norm(output)
return output
PyTorch 2.xではグラフモードが導入されましたが、依然として制御フローが発生すると、Pythonによる実行が行われます。
静的グラフ
静的グラフは、Pythonなどの高級言語で記述されたニューラルネットワークの構造とパラメータから固定された計算グラフを事前に生成する方式です。この方式は静的計算グラフとも呼ばれます。逆伝播の勾配計算は、計算グラフ上の最適化パスとして実装され、前方向計算グラフの損失関数を根ノードとして幅優先探索することで、対偶構造を用いて逆方向計算グラフが自動生成されます。
制御フロー原語
TensorFlowは制御フロー原語をサポートしており、制御フローの必要に応じて計算グラフに制御フロー原語を導入します。制御フロー原語は実行時に第一級の機能として扱われます。基本設計原則としては、少数の操作を含む原子的な演算子（Function型特殊演算子）を導入し、それらを組み合わせて複雑な制御フローを表現します。
TensorFlowでは、制御フローをサポートするための3層のアーキテクチャが採用されています。開発者向けのAPI、低レベルの制御フロー原語、そして計算グラフの最適化処理がそれぞれ対応しています。TensorFlowは抽象度の異なる複数のAPIと原語を提供し、最適化の機会を保ちつつプログラミングの利便性を高めています。
以下は、TensorFlow 2.xでの2重forループの例です：
i = tf.constant(0)
j = tf.constant(0)
a = lambda i: tf.less(i, 10)
b = lambda i: (tf.add(i, 1), )
c = lambda i: (tf.add(j, 1), )
y = tf.while_loop(a, b, [j])
r = tf.while_loop(c, y, [i])
TensorFlowでは、各演算子は実行フレーム内で実行され、各フレームはグローバルに一意な名前を持ちます。制御フロー原語はこれらのフレームを管理します。実行フレームは言語のスコープに似ており、演算子のコンテキスト情報を保持します。ループが導入されると、同じ演算子が複数回実行される可能性があり、実行フレームは実行時に作成されます。ネストされた実行フレームに対応し、開発者が記述した制御フロー構造を反映します。
制御フロー原語には以下の5つの原子演算子が含まれます：
- Switch：入力のブール値に従って入力テンソルを2つの出力に送信する。
- Merge：利用可能な入力のうち一つを出力に送信する。
- Enter：指定された名前で識別される実行フレームにテンソルを渡す。
- Exit：実行フレームから親フレームに戻す。
- NextIteration：現在の実行フレームの次のイテレーションにテンソルを渡す。
ソースコード解析
ソースコード解析方式では、フロントエンドで高級言語を解析して計算グラフに変換し、バックエンドで制御フロー構造をサブグラフとして展開します。forループなどの制御構造は、直接展開して順序付きの複数のサブグラフとして表現され、if-elseのような分岐はそれぞれのサブグラフを作成して、実行時に選択的に実行されます。
この方式の利点は：
- 開発者は制約付きでPythonの制御フロー構造を自由に使用できる。
- ホスト言語と実行プロセスを切り離すことで、実行効率が向上する。
- コンパイル時に全体の計算過程が把握できるため、実行時の最適化が可能。
ただし、静的グラフであるため、以下の欠点があります：
- ハードウェアがサポートしていない制御フローは、言語境界でのジャンプが発生し、実行時のオーバーヘッドが生じる。
- 一部の制御フロー構造は表現できないという制限がある。