torch.fx.GraphModule.graphプログラミング入門：基本と応用
2025-05-31
より詳しく見ていきましょう。
-
- FXは、PyTorchモデルを分析・変換するためのツールキットです。従来のTorchScriptがソースコードレベルでの変換を行うのに対し、FXはPyTorchの実行をトレースし、その操作をグラフ構造として捉えます。これにより、より柔軟で強力なモデルの解析や最適化が可能になります。
-
torch.fx.GraphModule
torch.fx.GraphModuleは、PyTorchのnn.ModuleをFXによって変換した結果得られるクラスです。これは、元のモジュールの計算グラフを保持しており、そのグラフを通じてモデルの構造や演算を操作できます。
-
graph属性
torch.fx.GraphModuleオブジェクトの持つ属性の一つがgraphです。このgraph属性は、torch.fx.Graph型のオブジェクトであり、モデルの計算グラフそのものを表しています。
このグラフ構造を利用することで、以下の様な操作が可能になります。
- コードの生成
変換されたグラフから、再びPyTorchのコード（PythonやTorchScript）を生成できます。 - グラフの変換
新しいノードの挿入、既存のノードの削除や置換など、グラフ構造を直接操作することで、モデルの最適化や特殊な処理の追加が行えます。 - モデルの構造分析
グラフ内のノードやエッジを辿ることで、モデルの各層の接続関係やデータの流れを把握できます。
例えば、あるnn.Moduleのインスタンス model があったとして、それをFXで変換すると以下のようになります。
import torch
import torch.nn as nn
import torch.fx
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 20)
self.relu = nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
return x
model = MyModule()
traced_model = torch.fx.symbolic_trace(model)
# traced_model は torch.fx.GraphModule のインスタンス
print(type(traced_model))
# <class 'torch.fx.graph_module.GraphModule'>
# traced_model.graph が計算グラフを表す torch.fx.Graph オブジェクト
graph = traced_model.graph
print(type(graph))
# <class 'torch.fx.graph.Graph'>
# グラフ内のノードをイテレートして情報を確認できます
for node in graph.nodes:
print(f"ノード名: {node.name}, オペコード: {node.op}, ターゲット: {node.target}, 引数: {node.args}, キーワード引数: {node.kwargs}")
グラフの不正な変更 (Graph Mutation Issues)
- トラブルシューティング
- graph.lint() の利用
グラフの整合性をチェックするためのメソッドgraph.lint()を実行し、エラーがないか確認します。 - graph.node_copy() や graph.erase_node() の慎重な使用
ノードのコピーや削除を行う際は、関連するエッジや他のノードへの影響を十分に考慮し、慎重に行います。 - 新しいノードの挿入方法の確認
新しいノードを挿入する際は、正しいop_code、target、args、kwargsを指定し、グラフに正しく接続されていることを確認します。graph.create_node()メソッドの利用が推奨されます。 - 変更後のグラフの可視化
変更後のグラフをtorch.fx.Graph.print_tabular()や、GraphVizなどのツールを使って可視化し、意図した通りの変更になっているか確認します。
- graph.lint() の利用
- エラー
グラフのノードやエッジを直接的かつ不適切に変更した場合、グラフの整合性が失われ、後続の処理で予期しないエラーが発生することがあります。例えば、存在しないノードを参照したり、ノードの入出力を矛盾した状態にしたりする場合などです。
シンボリック・トレース時のエラー (Symbolic Tracing Errors)
- トラブルシューティング
- トレース可能なコードの記述
モデルのforwardメソッド内で、FXがサポートする演算のみを使用するようにコードを修正します。リストや辞書の操作、ループの条件が入力テンソルに依存する場合などは、トレースが難しくなることがあります。 - torch.fx.wrap() の利用
トレースできない関数やモジュールをラップし、FXにブラックボックスとして扱うように指示できます。ただし、ラップされた部分はグラフの内部構造が解析できなくなります。 - torch.nn.Module の利用
複雑な処理は、可能な限りtorch.nn.Moduleのサブクラスとして実装し、その中でPyTorchの演算を使用するようにします。 - トレース結果の確認
生成されたグラフが意図した計算を表しているか、ノードの構成や接続を確認します。
- トレース可能なコードの記述
- エラー
torch.fx.symbolic_trace()関数でモデルをトレースする際に、Pythonの制御フロー（if文、for文など）や、FXが直接的に扱えない操作が含まれていると、トレースが中断されたり、不完全なグラフが生成されたりする可能性があります。
ノードのターゲットに関するエラー (Node Target Errors)
- トラブルシューティング
- ターゲットの確認
ノードのtarget属性が、意図した関数やメソッドを正しく参照しているか確認します。特に、カスタム関数やメソッドを使用している場合は、名前空間が正しいか注意が必要です。 - op_code と target の整合性
ノードのop_code（例: 'call_function', 'call_method', 'get_attr'）とtargetの種類が一致しているか確認します。例えば、op_codeが 'call_module' の場合は、targetはnn.Moduleの属性名である必要があります。
- ターゲットの確認
- エラー
各ノードが実行する演算のターゲット（通常は関数やメソッド名）が正しくない場合、グラフを実行しようとした際にエラーが発生します。
ノードの引数に関するエラー (Node Argument Errors)
- トラブルシューティング
- 引数の型と形状の確認
各ノードのargsおよびkwargsの内容を確認し、ターゲットの関数やメソッドの引数仕様と合致しているか確認します。 - 前のノードの出力の確認
現在のノードへの入力が、前のノードの出力として正しく接続されているか確認します。ノード間のデータの流れを意識することが重要です。
- 引数の型と形状の確認
- エラー
ノードへの入力引数 (argsおよびkwargs) が、ターゲットの関数やメソッドが期待する型や数と一致しない場合、グラフの実行時にエラーが発生します。
カスタム演算の扱い (Handling Custom Operations)
- トラブルシューティング
- torch.fx.wrap() の利用
前述の通り、トレースできないカスタム関数をラップしてFXに認識させることができます。 - torch.autograd.Function の利用
より複雑なカスタム演算を定義する場合は、torch.autograd.Functionを継承したクラスを作成し、そのforwardおよびbackwardメソッドを実装することで、PyTorchの自動微分機構と連携させることができます。FXでこれを扱う場合は、op_codeを 'call_function' とし、targetにカスタム関数のクラスを指定することが考えられます。
- torch.fx.wrap() の利用
- エラー
FXが標準でサポートしていないカスタムの関数や演算をモデル内で使用した場合、トレース時に正しくグラフ化されないことがあります。
グラフの実行時エラー (Graph Execution Errors)
- トラブルシューティング
- エラーメッセージの確認
Pythonの標準的なエラーメッセージを注意深く読み、どのノードでどのようなエラーが発生しているかを特定します。 - 中間出力の確認
必要に応じて、グラフの途中のノードの出力を検査するコードを挿入し、データの流れや値が期待通りであるか確認します。 - 個々のノードのテスト
問題が発生している可能性のあるノードに対応する演算を個別にテストし、入力に対して正しい出力が得られるか確認します。
- エラーメッセージの確認
- エラー
変更または生成したグラフをGraphModuleのインスタンスで実行しようとした際に、内部のノードの演算が失敗することがあります。
- 公式ドキュメントの参照
PyTorchの公式ドキュメントやFXに関する情報を参照し、正しいAPIの使用方法や注意点を確認します。 - PyTorchとFXのバージョンの確認
バージョンによって挙動が異なる場合があるため、使用しているPyTorchとFXのバージョンを確認します。 - 最小限の再現コード (Minimal Reproducible Example)
問題を報告する際や自身でデバッグを行う際は、できるだけ短いコードで問題を再現できるように努めます。
グラフ構造の確認
まず、FXでトレースしたモデルのグラフ構造を確認する基本的な例です。
import torch
import torch.nn as nn
import torch.fx
# 簡単なモデルを定義
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(10, 20)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(20, 5)
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.linear2(x)
return x
# モデルのインスタンスを作成し、FXでトレース
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
# GraphModule の graph 属性にアクセス
graph = traced_model.graph
# グラフ内のノードをイテレートして情報を表示
print("グラフ内のノード:")
for node in graph.nodes:
print(f" 名前: {node.name}, オペコード: {node.op}, ターゲット: {node.target}, 入力: {node.args}, キーワード引数: {node.kwargs}")
# グラフの構造をtabular形式で表示
print("\nグラフ構造 (tabular):")
graph.print_tabular()
このコードでは、SimpleModel を torch.fx.symbolic_trace() でトレースし、得られた traced_model の graph 属性にアクセスしています。その後、グラフ内の各ノードの基本的な情報（名前、オペコード、ターゲットなど）を表示し、さらに graph.print_tabular() を用いてグラフの構造をより分かりやすい表形式で出力しています。
新しいノードの追加
次に、既存のグラフに新しいノードを追加する例です。ここでは、ReLUの後に別のReLUを追加してみます。
import torch
import torch.nn as nn
import torch.fx
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 20)
self.relu = nn.ReLU()
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
return x
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
graph = traced_model.graph
# 'relu' ノードを探す
relu_node = None
for node in graph.nodes:
if node.op == 'call_module' and isinstance(traced_model.get_submodule(node.target), nn.ReLU):
relu_node = node
break
if relu_node:
# 新しい ReLU ノードを relu ノードの後に挿入
with graph.inserting_after(relu_node):
new_relu_node = graph.create_node(
op='call_module',
target='relu2', # 新しい submodule の名前
args=(relu_node,),
)
# 新しい ReLU submodule を GraphModule に追加
traced_model.relu2 = nn.ReLU()
# 元の出力ノードの入力を新しい ReLU ノードに変更
output_node = None
for node in graph.nodes:
if node.op == 'output':
output_node = node
break
if output_node:
output_node.replace_input_with(relu_node, new_relu_node)
# グラフを再コンパイル
traced_model.recompile()
# 変更後のグラフ構造を表示
print("変更後のグラフ構造 (tabular):")
traced_model.graph.print_tabular()
# 新しい GraphModule を使って推論を実行
input_tensor = torch.randn(1, 10)
output = traced_model(input_tensor)
print("\n出力:", output)
else:
print("ReLU ノードが見つかりませんでした。")
この例では、まずグラフ内の ReLU ノードを探し、その直後に新しい ReLU ノード (relu2) を挿入しています。graph.inserting_after() コンテキストマネージャーを使うと、指定したノードの後に新しいノードを挿入しやすくなります。また、新しい nn.Module のインスタンスを traced_model に追加し、元の出力ノードの入力を新しい ReLU ノードに変更しています。最後に traced_model.recompile() を呼び出して、変更されたグラフに基づいて実行可能なコードを再生成しています。
ノードの削除
次に、グラフから特定のノードを削除する例です。ここでは、最初の線形層 (linear1) を削除してみます（ただし、これを行うとモデルが壊れる可能性があります。あくまで例として見てください）。
import torch
import torch.nn as nn
import torch.fx
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(10, 20)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(20, 5)
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.linear2(x)
return x
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
graph = traced_model.graph
# 削除したい 'linear1' ノードを探す
linear1_node = None
for node in graph.nodes:
if node.op == 'call_module' and node.target == 'linear1':
linear1_node = node
break
if linear1_node:
# 'linear1' ノードの出力を、その後続の 'relu' ノードの入力にする
relu_node = None
for user in list(linear1_node.users): # users はイテレータなのでリストに変換して安全に操作
if user.op == 'call_module' and isinstance(traced_model.get_submodule(user.target), nn.ReLU):
relu_node = user
break
if relu_node:
relu_node.replace_input_with(linear1_node, list(linear1_node.args)[0]) # linear1 の入力を relu の入力にする
# 'linear1' ノードに関連する submodule を削除
del traced_model.linear1
# 'linear1' ノードをグラフから削除
graph.erase_node(linear1_node)
# グラフを再コンパイル
traced_model.recompile()
# 変更後のグラフ構造を表示
print("変更後のグラフ構造 (tabular):")
traced_model.graph.print_tabular()
# 注意: このモデルは 'linear1' が削除されているため、元の入力形状では動作しない可能性があります。
else:
print("linear1 ノードが見つかりませんでした。")
この例では、まず削除したい linear1 ノードを探し、その出力を後続の relu ノードの入力に変更しています。次に、traced_model から対応する submodule (linear1) を削除し、最後に graph.erase_node() を用いてグラフからノードを削除しています。削除後も traced_model.recompile() を忘れずに行う必要があります。
グラフの実行
torch.fx.GraphModule は、通常の nn.Module と同様に呼び出して実行できます。上記の例でグラフを変更した後も、traced_model(input_tensor) のように入力テンソルを与えて実行できます。
torch.fx.Transformer を利用したグラフ変換
torch.fx.Transformer は、グラフパターンマッチングと置換のメカニズムを提供し、より宣言的な方法でグラフを変換できます。特定のパターンに合致する部分グラフを見つけ出し、それを指定した新しい部分グラフで置き換えることができます。
import torch
import torch.nn as nn
import torch.fx
from torch.fx.passes.graph_transform import GraphModuleTransformation
# 変換ルールを定義する Transformer のサブクラス
class ReplaceReLUWithSigmoid(GraphModuleTransformation):
def pattern(self):
# マッチさせたい部分グラフのパターンを定義
class Pattern(nn.Module):
def __init__(self):
super().__init__()
self.relu = nn.ReLU()
def forward(self, x):
return self.relu(x)
return Pattern()
def replacement(self, inps, node):
# マッチしたノードを置き換える処理を定義
return [torch.sigmoid(inps[0])]
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 20)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(20, 5)
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
x = self.linear2(x)
return x
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
# Transformer のインスタンスを作成し、グラフに適用
transformer = ReplaceReLUWithSigmoid()
transformed_model = transformer(traced_model)
# 変換後のグラフ構造を表示
print("変換後のグラフ構造 (tabular):")
transformed_model.graph.print_tabular()
# 変換後のモデルで推論を実行
input_tensor = torch.randn(1, 10)
output = transformed_model(input_tensor)
print("\n出力:", output)
この例では、ReplaceReLUWithSigmoid という GraphModuleTransformation のサブクラスを定義しています。pattern() メソッドでReLUのパターンを指定し、replacement() メソッドでReLUノードをSigmoid関数で置き換える処理を記述しています。Transformer を使うと、グラフの構造を直接操作するよりも、意味的な変換をより簡潔に記述できます。
torch.fx.PassManager を利用した複数のグラフ変換のパイプライン
複数のグラフ変換を順番に適用したい場合、torch.fx.PassManager を利用できます。これにより、一連の変換処理をパイプラインとして管理し、実行することができます。
import torch
import torch.nn as nn
import torch.fx
from torch.fx.passes.graph_transform import GraphModuleTransformation
from torch.fx.passes.pass_manager import PassManager
# (上記の ReplaceReLUWithSigmoid クラスの定義は省略)
class AddBiasAfterLinear(GraphModuleTransformation):
def pattern(self):
class Pattern(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 20)
def forward(self, x):
return self.linear(x)
return Pattern()
def replacement(self, inps, node):
linear_module = self.get_module(node.target)
bias = torch.randn(linear_module.out_features)
return [torch.add(inps[0], bias)]
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(10, 20)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(20, 5)
def forward(self, x):
x = self.linear1(x)
x = self.relu(x)
x = self.linear2(x)
return x
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
# PassManager に適用する変換を登録
pm = PassManager([
ReplaceReLUWithSigmoid(),
AddBiasAfterLinear()
])
# グラフに変換パイプラインを適用
transformed_model = pm(traced_model)
# 変換後のグラフ構造を表示
print("変換後のグラフ構造 (tabular):")
transformed_model.graph.print_tabular()
# 変換後のモデルで推論を実行
input_tensor = torch.randn(1, 10)
output = transformed_model(input_tensor)
print("\n出力:", output)
この例では、ReplaceReLUWithSigmoid に加えて、線形層の後にバイアスを加える AddBiasAfterLinear という新しい変換を定義しています。PassManager にこれらの変換をリストとして渡し、トレースされたモデルに適用することで、複数の変換が順番に実行されます。
torch.fx.Node の属性を利用した間接的な操作
torch.fx.Node オブジェクトの属性（op, target, args, kwargs) を直接変更することも、グラフを操作する一つの方法ですが、前述の通り、グラフの整合性を保つためには注意が必要です。しかし、これらの属性を読み取り、分析に基づいて条件付きでノードを追加・削除するなどの処理を行うことができます。
例えば、特定の属性を持つノードを検索し、その情報に基づいて新しいノードを挿入したり、既存のノードの引数を変更したりする処理などが考えられます。
カスタムのグラフ解析パスの作成
より複雑な分析や変換を行うために、torch.fx.Graph を直接操作するカスタムのパス（pass）を作成することもできます。これは、グラフ内のノードをイテレートし、特定の条件を満たすノードに対して何らかの処理を行う関数として実装できます。
import torch
import torch.nn as nn
import torch.fx
def fuse_relu_linear(graph: torch.fx.Graph, model: torch.fx.GraphModule):
nodes_to_remove = set()
for node in graph.nodes:
if node.op == 'call_module' and isinstance(model.get_submodule(node.target), nn.ReLU):
for user in list(node.users):
if user.op == 'call_module' and isinstance(model.get_submodule(user.target), nn.Linear):
# ReLU と Linear を融合する処理 (ここでは簡略化)
print(f"ReLU {node.name} と Linear {user.name} を融合します (実際にはより複雑な処理が必要)")
nodes_to_remove.add(node)
# ... 実際の融合ロジック ...
for node in nodes_to_remove:
graph.erase_node(node)
graph.lint()
model.recompile()
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 20)
self.relu = nn.ReLU()
self.linear2 = nn.Linear(20, 5)
def forward(self, x):
x = self.linear(x)
x = self.relu(x)
x = self.linear2(x)
return x
model = SimpleModel()
traced_model = torch.fx.symbolic_trace(model)
# カスタムのグラフ解析パスを適用
fuse_relu_linear(traced_model.graph, traced_model)
# 変換後のグラフ構造を表示
print("変換後のグラフ構造 (tabular):")
traced_model.graph.print_tabular()
この例は、ReLU の直後に Linear 層が続くパターンを見つけ、それらを融合しようとするカスタムのパス fuse_relu_linear を示しています（実際の融合ロジックは簡略化されています）。このようなカスタムパスを作成することで、特定の最適化や分析を柔軟に行うことができます。