PyTorch FXグラフ解析の決定版：print_tabular()で見る内部構造
torch.fxは、PyTorchのプログラムをシンボリックにトレースし、Pythonコードを表現するグラフ（torch.fx.Graphオブジェクト）を構築するためのモジュールです。このグラフは、モデルの構造や各操作の流れを詳細に把握するために非常に役立ちます。
torch.fx.Graph.print_tabular()メソッドは、このtorch.fx.Graphオブジェクトの内容を**表形式（tabular format）**で整形して表示するためのユーティリティ関数です。グラフ内の各ノード（操作、関数呼び出し、定数など）に関する情報を、人間が読みやすい形式で出力してくれます。
なぜ「print_tabular()」が便利なのか？
- 教育と理解
PyTorchの内部動作やFXグラフの仕組みを学習する上で、具体的なグラフの構造を視覚的に確認できることは、理解を深める助けになります。 - デバッグと解析
モデルの変換（例: 量子化、グラフ最適化）や、特定の操作が意図通りにトレースされているかを確認する際に、各ノードの入力、出力、ターゲットなどの詳細を素早く把握できます。 - 視認性の向上
torch.fx.Graphオブジェクトをそのまま表示すると、Pythonオブジェクトの内部表現が表示されるため、非常に読みにくい場合があります。print_tabular()は、情報を整理された表形式で提示するため、一目でグラフの構造を理解しやすくなります。
出力される主な情報
print_tabular()によって出力される表には、通常、以下の情報が含まれます。
kwargs: ノードのキーワード引数。args: ノードの引数。target: ノードが表す具体的な関数、モジュール、または属性。name: ノードの名前。通常、Pythonコードの変数名などから自動的に生成されます。opcode: ノードの種類（例:placeholder,call_function,call_module,call_method,get_attr,output）。placeholder: グラフへの入力。call_function: Pythonの関数呼び出し（例:torch.add,F.relu）。call_module:torch.nn.Moduleのサブモジュールの呼び出し（例:self.conv1）。call_method: オブジェクトのメソッド呼び出し（例:x.view）。get_attr: 属性の取得（例:self.param）。output: グラフの出力。
使用例
簡単なPyTorchモデルを例に、print_tabular()の使い方を見てみましょう。
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 5)
def forward(self, x):
x = self.linear(x)
x = torch.relu(x)
return x
# モデルのインスタンス化
model = SimpleModel()
# モデルをシンボリックにトレース
traced_model = symbolic_trace(model)
# グラフを取得
graph = traced_model.graph
# グラフを表形式で表示
print("--- Graph Tabular View ---")
graph.print_tabular()
print("------------------------")
# グラフのコードも確認（参考）
print("\n--- Graph Python Code ---")
print(graph.python_code)
print("-------------------------")
このコードを実行すると、print_tabular()の出力として、以下のような表が表示されます（実際の出力は環境やPyTorchのバージョンによって若干異なる場合があります）。
--- Graph Tabular View ---
opcode name target args kwargs
------------- ------- --------------------- --------------- --------
placeholder x x () {}
call_module linear linear (x,) {}
call_function relu <built-in method relu> (linear,) {}
output output output (relu,) {}
------------------------
--- Graph Python Code ---
def forward(self, x):
linear = self.linear(x)
relu = torch.relu(linear)
return relu
-------------------------
上記の表から、入力xがlinearモジュールに渡され、その出力がtorch.relu関数に渡され、最後にreluの出力がグラフの出力となる一連の流れが明確に読み取れます。
tabulateライブラリが見つからない (ModuleNotFoundError)
これは最も一般的なエラーの一つです。print_tabular()は、テーブル表示のために外部ライブラリであるtabulateに依存しています。もしこのライブラリがインストールされていない場合、ModuleNotFoundErrorが発生します。
エラーメッセージ例
ModuleNotFoundError: No module named 'tabulate'
トラブルシューティング
tabulateライブラリをインストールするだけです。
pip install tabulate
FXがグラフを正しくトレースできない (TraceError / グラフブレイク)
torch.fx.symbolic_traceは、Pythonのコードをシンボリックに解析してグラフを構築します。しかし、全てのPythonコードがトレース可能であるわけではありません。トレースが失敗したり、意図しない形でグラフが「ブレイク」したりすることがあります。この場合、print_tabular()自体はエラーを出さないかもしれませんが、出力されるグラフが期待と異なるか、非常に短くなってしまいます。
主な原因とトラブルシューティング
-
カスタムクラスやカスタム演算子
独自のクラスや演算子を定義している場合、FXがそれらをどのようにトレースするかを明示的に指定する必要がある場合があります。トラブルシューティング
torch.fx.wrapを使用して、特定の関数をFXグラフの単一のノードとして扱うように指定できます。また、カスタム演算子の場合は、PyTorchのオペレーター拡張メカニズム（torch.libraryなど）を適切に使用する必要があります。 -
タプルやリストの*args / **kwargs展開
Proxyオブジェクトをループ内で使用したり、*argsや**kwargsとして関数引数に渡したりすると、torch.fx.proxy.TraceError: Proxy object cannot be iterated.のようなエラーが発生することがあります。エラーメッセージ例
torch.fx.proxy.TraceError: Proxy object cannot be iterated.トラブルシューティング
FXグラフのノードは、中間的なProxyオブジェクトとして扱われます。これらのProxyオブジェクトは、通常のPythonオブジェクトのようにイテレートしたり、動的にアクセスしたりすることはできません。イテレーションや動的なアクセスが必要な場合は、グラフ構築後にこれらの操作を行うようにコードを再構築するか、FXのトレースの限界を理解し、トレースできない部分を明示的にPyTorchのEagerモードにフォールバックさせる（torch.compileのグラフブレイクのように）必要があります。 -
アサーション (assert) の使用
Pythonの通常のassert文はトレース中にエラーを引き起こすことがあります。トラブルシューティング
代わりにtorch._assert（PyTorch内部で使用されるトレーサブルなアサーション）を使用することを検討します。ただし、これはプライベートAPIなので、注意が必要です。 -
インプレース操作の複雑さ
テンソルのスライスに対するインプレース変更（例:x[:, 0] += 1）は、トレースが難しい場合があります。トラブルシューティング
新しい変数にスライスをコピーし、操作を適用した後で元のテンソルを再構築するなど、インプレースではない方法で表現することを検討します。 -
サポートされていないPythonの組み込み関数やモジュール
FXはPyTorchの操作に焦点を当てており、全てのPythonの組み込み関数や標準ライブラリのモジュールがトレースできるわけではありません。例えば、inspectモジュールの一部などが挙げられます。トラブルシューティング
可能であれば、トレース中にそれらの関数を呼び出さないようにコードを修正します。 -
データ依存の制御フロー (Data-dependent control flow)
入力テンソルの値に依存するif文やforループは、FXでは静的なグラフとして表現できません。def forward(self, x): if x.sum() > 0: # データ依存の制御フロー return x + 1 else: return x - 1トラブルシューティング
可能な限り、データ依存の制御フローを避けるようにモデルを再構築します。PyTorchのテンソル演算（例:torch.where）で表現できる場合は、それを使用します。def forward(self, x): condition = x.sum() > 0 return torch.where(condition, x + 1, x - 1)完全に避けられない場合は、
torch.compileのようなより高度なトレーシングツールを検討するか、FXの実験的な機能（ASTリライターなど）を試す必要があります。
print_tabular()自体はエラーを出さないものの、出力されるグラフが期待していたものと異なる場合があります。これは、トレースが部分的にしか成功していないか、モデルの構造が思っていたものと違う場合に発生します。
トラブルシューティング
-
torch.compileとの連携
PyTorch 2.0以降では、torch.compileがFXのバックエンドとしてTorchDynamoを使用しています。torch.compileは、従来のsymbolic_traceよりも多くのPythonコードパターンを扱え、自動的にグラフブレイクを処理してくれます。もしsymbolic_traceで問題が発生する場合は、まずtorch.compileを試してみることを強くお勧めします。torch.compileを使用した場合でも、バックエンド内でgm.graph.print_tabular()を呼び出すことでグラフの構造を確認できます。import torch import torch.nn as nn from torch.fx import symbolic_trace class MyModel(nn.Module): def forward(self, x): # 例: データ依存の制御フロー if x.mean() > 0: x = x + 1 else: x = x - 1 return x * 2 model = MyModel() # torch.compile を使用してトレース # TorchDynamo がグラフブレイクを処理する compiled_model = torch.compile(model) # グラフの確認（compileの内部でprint_tabularを呼び出すように設定することも可能） # この例では直接print_tabularは呼ばれませんが、デバッグモードで確認できます x = torch.randn(5) _ = compiled_model(x) # 実際に一度実行してコンパイルをトリガーTORCH_COMPILE_DEBUG=1のような環境変数を設定してtorch.compileを実行すると、より詳細なデバッグ情報（グラフブレイクの原因など）が出力されることがあります。 -
サブモジュールの確認
nn.Sequentialやカスタムのnn.Moduleで構成されたモデルの場合、print_tabular()の出力はトップレベルのモジュールとその直接の子モジュールを示します。もし、深いネストされたサブモジュールの内部構造まで確認したい場合は、そのサブモジュールを個別にトレースする必要があります。 -
モデルのforwardメソッドの確認
torch.fx.symbolic_traceは、デフォルトでnn.Moduleのforwardメソッドをトレースします。もし、トレースしたいロジックが別のメソッドにある場合、明示的にそのメソッドをトレースする必要があります。
torch.fxは、PyTorchモデルのシンボリックなトレースを行い、計算グラフを構築するための強力なツールです。print_tabular()はこの構築されたグラフを人間が読みやすい表形式で表示するのに役立ちます。
以下の例では、様々なシナリオでprint_tabular()がどのように使われ、どのような出力が得られるかを説明します。
例1: 基本的な線形モデルのトレース
最もシンプルなケースとして、nn.Linearとtorch.reluを含むモデルをトレースします。
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
# 1. シンプルなPyTorchモデルを定義
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 5) # 入力10、出力5の線形層
def forward(self, x):
x = self.linear(x) # 線形変換
x = torch.relu(x) # ReLU活性化関数
return x
# 2. モデルのインスタンスを作成
model = SimpleModel()
# 3. モデルをシンボリックにトレース
# symbolic_traceは、モデルのforwardメソッドを解析し、計算グラフを作成します。
# 戻り値はGraphModuleという特殊なnn.Moduleで、内部にグラフ情報を持っています。
traced_model = symbolic_trace(model)
# 4. トレースされたモデルからGraphオブジェクトを取得
# GraphModuleの.graph属性にGraphオブジェクトが格納されています。
graph = traced_model.graph
# 5. Graphオブジェクトを表形式で表示
print("--- 例1: SimpleModel のグラフ表示 ---")
graph.print_tabular()
print("------------------------------------")
# 参考: グラフに対応するPythonコードも表示
print("\n--- 例1: SimpleModel のPythonコード ---")
print(graph.python_code)
print("--------------------------------------")
コードの解説
SimpleModelという簡単なニューラルネットワークを定義します。このモデルは、nn.Linear層とtorch.relu関数から構成されています。SimpleModelのインスタンスを作成します。torch.fx.symbolic_trace(model)を使って、モデルのforwardメソッドの実行フローを解析し、torch.fx.GraphModuleオブジェクトに変換します。このGraphModuleが、モデルの計算グラフを抽象的に表現しています。traced_model.graphから、実際の計算グラフを表すtorch.fx.Graphオブジェクトを取り出します。graph.print_tabular()を呼び出すことで、この計算グラフの各ノード（操作）が表形式で整形されて出力されます。
--- 例1: SimpleModel のグラフ表示 ---
opcode name target args kwargs
------------- ------- --------------------- --------------- --------
placeholder x x () {}
call_module linear linear (x,) {}
call_function relu <built-in method relu> (linear,) {}
output output output (relu,) {}
------------------------------------
--- 例1: SimpleModel のPythonコード ---
def forward(self, x):
linear = self.linear(x)
relu = torch.relu(linear)
return relu
--------------------------------------
- output
グラフの最終出力。 - call_function (relu)
torch.reluというPythonの関数が呼び出されていることを示します。targetは組み込みのreluメソッドです。 - call_module (linear)
self.linearモジュールが呼び出されていることを示します。targetはlinearという名前のモジュールです。 - placeholder (x)
グラフへの入力（ここではx）。
この表は、xがlinearに入力され、その結果がreluに入力され、最終的にreluの出力が返されるという、モデルの計算フローを明確に示しています。
例2: 複数のサブモジュールと定数を持つモデルのトレース
もう少し複雑なモデルで、複数のレイヤーや定数、異なるタイプのノードがどのように表示されるかを確認します。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.fx import symbolic_trace
# 1. 少し複雑なモデルを定義
class ComplexModel(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.bn1 = nn.BatchNorm2d(16)
self.pool = nn.MaxPool2d(2, 2)
self.const_add = 10.0 # 定数
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = F.relu(x) # torch.nn.functional を使用
x = self.pool(x)
x = x + self.const_add # 定数との加算
return x
# 2. モデルのインスタンスを作成
model = ComplexModel()
# 3. ダミー入力テンソル（トレースには型情報が必要なため）
dummy_input = torch.randn(1, 3, 32, 32) # (バッチサイズ, チャンネル, 高さ, 幅)
# 4. モデルをシンボリックにトレース
# 入力テンソルを渡すことで、テンソルの形状や型情報がトレースに利用されます。
traced_model = symbolic_trace(model, concrete_args={'x': dummy_input})
# 5. Graphオブジェクトを取得
graph = traced_model.graph
# 6. Graphオブジェクトを表形式で表示
print("\n--- 例2: ComplexModel のグラフ表示 ---")
graph.print_tabular()
print("---------------------------------------")
コードの解説
ComplexModelは、畳み込み層、バッチ正規化層、ReLU（F.reluを使用）、Maxプーリング層、そして定数との加算を含みます。dummy_inputは、トレース時にモデルが受け取る入力の形状と型をFXに伝えるために使用されます。symbolic_traceにconcrete_argsとして渡します。- トレース後、
graph.print_tabular()によってグラフが表示されます。
--- 例2: ComplexModel のグラフ表示 ---
opcode name target args kwargs
------------- ------------ -------------------------- ---------------------------------------------- --------
placeholder x x () {}
call_module conv1 conv1 (x,) {}
call_module bn1 bn1 (conv1,) {}
call_function relu <function relu at 0x...> (bn1,) {} # F.relu
call_module pool pool (relu,) {}
get_attr const_add const_add () {} # モデルの属性
call_function add <built-in method add> (pool, const_add) {} # 加算
output output output (add,) {}
---------------------------------------
- call_function (add)
テンソルの加算（x + self.const_add）が、内部的にはtorch.addのような関数呼び出しとしてトレースされていることを示します。 - get_attr (const_add)
モデルのインスタンス変数（属性）であるself.const_addがアクセスされていることを示します。 - call_function (relu)
torch.nn.functional.reluが呼び出されています。targetには関数の参照が表示されます。 - call_module
conv1,bn1,poolといったnn.Moduleのインスタンスが呼び出されていることを示します。
この例から、FXがnn.Module、torch.nn.functionalの関数、およびモデルの属性アクセスをどのように異なる種類のノードとして表現するかが分かります。
例3: トレース不可能なパターンとprint_tabular()の限界
torch.fxは全てのPythonコードをトレースできるわけではありません。特に、データ依存の制御フロー（if文やforループがテンソルの値に依存する場合）はトレースが困難です。このような場合、print_tabular()は期待通りのグラフを表示しないか、エラーを出すことがあります。
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
# 1. トレースが困難なパターンを含むモデル
class ProblematicModel(nn.Module):
def __init__(self):
super().__init__()
self.param = nn.Parameter(torch.tensor(0.5))
def forward(self, x):
# データ依存の条件分岐
# x.sum() の結果によってパスが変わるため、静的なグラフに変換できない
if x.sum() > self.param: # この行が問題
return x * 2
else:
return x / 2
# 2. モデルのインスタンスとダミー入力
model = ProblematicModel()
dummy_input = torch.randn(5)
# 3. モデルをシンボリックにトレース
# このトレースはエラーを出す可能性があります (TraceError)
# あるいは、if文の条件がFalseと評価され、片方のパスしかトレースされないことがあります。
try:
print("\n--- 例3: ProblematicModel のグラフ表示 (試行) ---")
traced_model = symbolic_trace(model, concrete_args={'x': dummy_input})
graph = traced_model.graph
graph.print_tabular()
except Exception as e:
print(f"トレース中にエラーが発生しました: {e}")
print("--------------------------------------------------")
# 解決策のヒント: torch.where を使用
class ResolvedModel(nn.Module):
def __init__(self):
super().__init__()
self.param = nn.Parameter(torch.tensor(0.5))
def forward(self, x):
# データ依存の条件分岐を torch.where で表現
condition = x.sum() > self.param
return torch.where(condition, x * 2, x / 2)
model_resolved = ResolvedModel()
traced_model_resolved = symbolic_trace(model_resolved, concrete_args={'x': dummy_input})
graph_resolved = traced_model_resolved.graph
print("\n--- 例3: ResolvedModel のグラフ表示 (解決済み) ---")
graph_resolved.print_tabular()
print("--------------------------------------------------")
コードの解説
ProblematicModelでは、x.sum() > self.paramという条件がテンソルの実際の値に依存するため、FXはグラフを静的に構築できません。この場合、TraceErrorが発生するか、あるいはsymbolic_traceがdummy_inputの値に基づいて一方のパスしかトレースしない可能性があります。try-exceptブロックでエラーを捕捉し、問題が発生したことを示します。ResolvedModelでは、torch.whereを使ってデータ依存の条件分岐をPyTorchの操作に変換しています。torch.whereはテンソル操作なので、FXによって正しくトレースできます。
ProblematicModelの部分では、TraceErrorや「トレース中にエラーが発生しました」のようなメッセージが表示されます。
ResolvedModelの部分では、以下のような表が表示されます。
--- 例3: ResolvedModel のグラフ表示 (解決済み) ---
opcode name target args kwargs
------------- --------- ---------------------------- ---------------------------------------------- --------
placeholder x x () {}
get_attr param param () {}
call_function sum <built-in method sum> (x,) {}
call_function gt <built-in method gt> (sum, param) {} # 比較演算子 >
call_function mul <built-in method mul> (x, 2) {} # x * 2
call_function truediv <built-in method truediv> (x, 2) {} # x / 2
call_function where <function where at 0x...> (gt, mul, truediv) {} # torch.where
output output output (where,) {}
--------------------------------------------------
where:torch.where関数が正しくトレースされ、条件（gt）、真の場合のパス（mul）、偽の場合のパス（truediv）を引数として受け取っていることが分かります。mul,truediv: 乗算と除算がそれぞれトレースされます。gt: 比較演算子>がtorch.gtのような内部関数としてトレースされます。
この例は、torch.fxのトレースの限界と、それを克服するための一般的なパターン（torch.whereの使用）を示しており、print_tabular()がデバッグにどのように役立つかを強調しています。
torch.fx.Graph.print_tabular()は、FX (Functional Transformations) が生成した計算グラフを整形して表示するのに非常に便利ですが、PyTorchのグラフ構造を理解したり、デバッグしたりするための他の方法もいくつか存在します。これらの代替方法には、それぞれ異なるユースケースと利点があります。
主な代替方法を以下に示します。
- torch.fx.Graph.python_code の利用
- torch.fx.Graph.nodes を直接イテレート
- Graphviz を使用したグラフの可視化
- torch.compile (TorchDynamo) のデバッグ機能
- サードパーティ製ツールや拡張機能
torch.fx.Graph.python_code の利用
print_tabular()が表形式でノード情報を表示するのに対し、graph.python_code属性は、FXグラフを再構築可能なPythonコードの文字列として提供します。これは、グラフがどのようにPyTorchの操作に変換されたかを直接的に理解するのに役立ちます。
利点
print_tabular()よりも、論理的なフローが把握しやすい場合がある。- 必要に応じて、このコードをコピーして実行し、さらにデバッグできる。
- グラフが生成された元のPythonコードに近い形で構造を理解できる。
欠点
- 表形式のような整理された列がないため、特定の情報（
opcodeやtargetなど）を見つけるのが難しい場合がある。 - 大規模なグラフではコードが長くなりすぎる可能性がある。
使用例
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 5)
def forward(self, x):
x = self.linear(x)
x = torch.relu(x)
return x
model = SimpleModel()
traced_model = symbolic_trace(model)
graph = traced_model.graph
print("--- 1. graph.python_code の利用 ---")
print(graph.python_code)
print("---------------------------------")
出力例
--- 1. graph.python_code の利用 ---
def forward(self, x):
linear = self.linear(x)
relu = torch.relu(linear)
return relu
---------------------------------
torch.fx.Graph.nodes を直接イテレート
torch.fx.Graphオブジェクトは、ノード（torch.fx.Nodeオブジェクト）のリストを内部に持っています。このリストを直接イテレートすることで、各ノードの詳細な属性にプログラム的にアクセスできます。これにより、特定の条件に基づいてノードをフィルタリングしたり、カスタムの表示形式を実装したりすることが可能です。
利点
- カスタムのレポートやデバッグツールを構築できる。
- 特定のノードタイプや引数を持つノードを検索できる。
- プログラムによる詳細な制御と解析が可能。
欠点
- 自分で出力フォーマットを実装する必要がある。
print_tabular()のような、すぐに読みやすい整形済み出力は得られない。
使用例
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(3, 16, 3)
self.pool = nn.MaxPool2d(2, 2)
def forward(self, x):
x = self.conv(x)
x = self.pool(x)
return x
model = SimpleModel()
dummy_input = torch.randn(1, 3, 32, 32)
traced_model = symbolic_trace(model, concrete_args={'x': dummy_input})
graph = traced_model.graph
print("\n--- 2. graph.nodes を直接イテレート ---")
for node in graph.nodes:
print(f" Node Name: {node.name}")
print(f" Opcode: {node.op}") # ノードの種類 (placeholder, call_module, call_functionなど)
print(f" Target: {node.target}") # ターゲットとなる関数、モジュール、属性
print(f" Args: {node.args}") # 位置引数
print(f" Kwargs: {node.kwargs}") # キーワード引数
print(f" Users: {[user.name for user in node.users]}") # このノードの出力を利用するノード
print("-" * 30)
print("---------------------------------------")
出力例
--- 2. graph.nodes を直接イテレート ---
Node Name: x
Opcode: placeholder
Target: x
Args: ()
Kwargs: {}
Users: ['conv']
------------------------------
Node Name: conv
Opcode: call_module
Target: conv
Args: (x,)
Kwargs: {}
Users: ['pool']
------------------------------
Node Name: pool
Opcode: call_module
Target: pool
Args: (conv,)
Kwargs: {}
Users: ['output']
------------------------------
Node Name: output
Opcode: output
Target: output
Args: (pool,)
Kwargs: {}
Users: []
------------------------------
---------------------------------------
Graphviz を使用したグラフの可視化
print_tabular()はテキストベースの表ですが、Graphvizのようなツールと連携させることで、FXグラフを視覚的に表現できます。これは、特に複雑なモデルや、ノード間の依存関係を直感的に把握したい場合に非常に有効です。
torch.fx自体はGraphvizの直接的なエクスポート機能を持っていませんが、torchvizなどの外部ライブラリや、FXグラフ情報をGraphvizのDOT言語に変換するカスタムスクリプトを作成することで実現できます。
利点
- 大規模なモデルでも、フローが明確になる。
- グラフの全体構造とノード間の接続を視覚的に把握できる。
欠点
- セットアップに手間がかかる場合がある。
- Graphvizと関連ツール（例:
graphvizPythonパッケージ）のインストールが必要。
使用例 (概念的なもの、Graphvizパッケージが必要)
# 例: torchviz を使用した（より一般的な）グラフ可視化の概念
# pip install graphviz torchviz
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
# from torchviz import make_dot # torchviz は FX グラフの直接可視化には向かない場合があります
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 5)
def forward(self, x):
x = self.linear(x)
x = torch.relu(x)
return x
model = SimpleModel()
dummy_input = torch.randn(1, 10)
traced_model = symbolic_trace(model, concrete_args={'x': dummy_input})
graph = traced_model.graph
print("\n--- 3. Graphviz によるグラフ可視化 (概念) ---")
print("FXグラフをGraphvizのDOT形式に変換し、画像として保存するカスタムスクリプトやツールが必要です。")
print("例: https://pytorch.org/docs/stable/fx.html#a-simple-use-case") # FXドキュメントの例を参照
# 以下は torchviz の使用例ですが、これは FX Graph ではなく、Eager モードの計算グラフを可視化します。
# f_graph = make_dot(traced_model(dummy_input), params=dict(traced_model.named_parameters()))
# f_graph.render("fx_graph_viz", format="png")
# print("生成されたグラフ: fx_graph_viz.png")
# FX Graph を直接 Graphviz にエクスポートするサンプルコード（より複雑）:
# 以下は概念的なもので、実際の実装には additional code が必要です。
# dot_graph = 'digraph G {\n'
# for node in graph.nodes:
# dot_graph += f' "{node.name}" [label="{node.op}\\n{node.target}"];\n'
# for user_node in node.users:
# dot_graph += f' "{node.name}" -> "{user_node.name}";\n'
# dot_graph += '}\n'
# with open("fx_graph.dot", "w") as f:
# f.write(dot_graph)
# print("DOTファイルが生成されました: fx_graph.dot (Graphvizでレンダリング可能)")
print("------------------------------------------")
torch.compile (TorchDynamo) のデバッグ機能
PyTorch 2.0 以降で導入されたtorch.compileは、内部的にFXグラフを生成するTorchDynamoを使用しています。torch.compileは、多くのケースでsymbolic_traceよりもロバストにグラフを生成できます。torch.compileは、デバッグのためにFXグラフの表示をサポートしています。
利点
print_tabular()よりも高度なグラフ最適化と実行を行う。- コンパイルプロセスのデバッグ情報としてFXグラフを出力できる。
- より広範なPythonコードをトレース可能（グラフブレイクを自動処理）。
欠点
- 直接
print_tabular()を呼び出すのではなく、torch.compileのデバッグフラグを設定する必要がある。
使用例
import torch
import torch.nn as nn
class MyComplexModel(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(10, 20)
self.fc2 = nn.Linear(20, 5)
def forward(self, x):
x = self.fc1(x)
x = torch.relu(x)
# 条件分岐を含む（TorchDynamoがこれを処理）
if x.sum() > 0:
x = x + 1
else:
x = x - 1
x = self.fc2(x)
return x
model = MyComplexModel()
dummy_input = torch.randn(1, 10)
print("\n--- 4. torch.compile のデバッグ機能 ---")
# 環境変数でデバッグモードを有効にする
# この設定はスクリプト実行前にコマンドラインで行うのが一般的ですが、
# Pythonコード内で一時的に設定することも可能です。
# 例: os.environ["TORCH_COMPILE_DEBUG"] = "1"
# または、Python 2.3 以降:
# torch._dynamo.config.log_level = 'DEBUG'
# torch._dynamo.config.debug = True
# 最も簡単な方法は、logging を設定することです
import logging
logging.basicConfig(level=logging.DEBUG) # もっとも詳細なログレベル
# compile() の backends.debug_backend を使用すると、グラフを見やすい形で出力できます
# (ただし、これはデバッグ用であり、実際のパフォーマンスには向きません)
# コンパイルがトリガーされた時に、デバッグ情報としてFXグラフが出力されることがあります。
# 例えば、`TORCH_COMPILE_DEBUG=1` 環境変数を設定して実行すると、
# コンソールにグラフ情報（tabular形式を含む）が出力されることがあります。
# または、`torch._dynamo.export()` を使用して、明示的にグラフをエクスポートする。
# traced_model_dynamo = torch._dynamo.export(model, dummy_input)
# print(traced_model_dynamo.graph.print_tabular()) # このように直接アクセスできる場合もある
# もっと簡単なデバッグ方法として、以下のようにバックエンドを指定します
try:
compiled_model = torch.compile(model, backend="eager") # eagerは最適化せず、グラフ生成のみを行う
# 実際には、この実行中にデバッグログが出力されます
_ = compiled_model(dummy_input)
print("torch.compile を使用してモデルがコンパイルされ、実行されました。")
print("デバッグログ (FXグラフ情報を含む) が出力されている可能性があります。")
except Exception as e:
print(f"torch.compile の実行中にエラーが発生しました: {e}")
print("---------------------------------------------")
サードパーティ製ツールや拡張機能
PyTorchのエコシステムには、計算グラフの可視化や解析を専門とするサードパーティ製ツールも存在します。これらは、torch.fxの機能に加えて、より高度な機能や異なる視覚化オプションを提供する場合があります。
利点
- 特定の解析タスクに特化している場合がある。
- よりリッチなGUIやインタラクティブな可視化機能。
欠点
torch.fxの最新の変更に常に対応しているとは限らない。- 追加のインストールと学習が必要。
例
- カスタムのグラフパーサー/ビジュアライザー
複雑な分析のために、torch.fx.Graph.nodesをイテレートし、独自のデータ構造や可視化ツールに変換するスクリプトを開発することも可能です。 - Netron
ONNX形式のモデルを可視化するツールですが、PyTorchモデルをONNXにエクスポートしてからNetronで見ることで、計算グラフを確認できます。
torch.fx.Graph.print_tabular()は手軽にFXグラフを確認できる便利な方法ですが、以下のような代替手段も検討できます。
- 専門的な可視化や解析機能が必要な場合
サードパーティ製ツールやカスタム実装。 - より複雑なモデルのトレースや最適化のデバッグが必要な場合
torch.compileのデバッグ機能の活用。 - グラフの全体構造を視覚的に把握したい場合
Graphviz (カスタムスクリプトまたは関連ライブラリ経由) を使用した可視化。 - 詳細なノード情報やプログラマブルなアクセスが必要な場合
graph.python_codeやgraph.nodesを直接イテレートする方法。