PyTorch FX不要コード削除(eliminate_dead_code)徹底解説と代替手法

torch.fx.Graph.eliminate_dead_code() は、PyTorch FX グラフ内の不要なコード（dead code）を削除するためのメソッドです。FX グラフは、PyTorch モデルを関数型の中間表現として表現したもので、最適化や変換といった処理を行うために用いられます。

具体的に「不要なコード」とは、グラフの出力に影響を与えない演算（ノード） のことを指します。これらの演算は、計算結果がどこにも利用されないため、実行しても意味がありません。eliminate_dead_code() を呼び出すことで、このような無駄な演算がグラフから取り除かれ、グラフがより簡潔で効率的になります。

このメソッドは、以下のような場合に役立ちます。

• モデルの軽量化
  不要な演算を削除することで、最終的なモデルのサイズや推論時の計算量をわずかに削減できる可能性があります。
• 手動でのグラフ編集後
  FX グラフを手動で操作してノードを追加・削除した場合、意図せず不要なノードが残ってしまう可能性があります。このメソッドでクリーンアップできます。
• グラフ変換後の整理
  グラフに対して様々な変換（例えば、演算の融合など）を行った後に、不要になった中間的な演算が残ることがあります。eliminate_dead_code() を適用することで、これらの残骸を取り除くことができます。

使用方法の例：

import torch
import torch.fx

# 簡単なモデルの定義
class MyModule(torch.nn.Module):
    def forward(self, x):
        y = x + 1
        z = x * 2
        return y

# モデルから FX グラフを作成
model = MyModule()
graph = torch.fx.symbolic_trace(model)

# グラフのノードを確認（削除前）
print("削除前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# 不要なコードを削除
graph.eliminate_dead_code()

# 削除後のグラフを再コンパイル
graph.lint()
new_module = torch.fx.GraphModule(model, graph)

# グラフのノードを確認（削除後）
print("\n削除後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# 新しいモジュールを使用
input_tensor = torch.randn(5)
output = new_module(input_tensor)
print("\n出力:", output)

上記の例では、MyModule の forward メソッドで z = x * 2 という演算を行っていますが、その結果 z は最終的な出力には使われていません。eliminate_dead_code() を呼び出すことで、この mul 演算（z の計算）がグラフから削除されることが期待されます。

• 削除されたノードに対応するパラメータなどは、元の torch.nn.Module からは削除されません。
• 削除後、グラフの整合性を保つために graph.lint() を呼び出すことが推奨されます。
• eliminate_dead_code() は、グラフの構造を解析し、出力に影響を与えないノードを特定して削除します。

---

意図しないノードの削除 (Unexpected Node Removal)

• トラブルシューティング
  
  • グラフの確認
    eliminate_dead_code() を実行する前後のグラフを注意深く比較し、意図せず削除されたノードがないか確認します。ノードの名前 (node.name) や演算の種類 (node.op)、入力 (node.all_input_nodes)、出力 (node.outputs) を調べます。
  • 出力ノードの確認
    グラフの出力ノード (graph.output) が正しく設定されているか確認します。出力ノードに繋がっていないパス上のノードは削除対象となる可能性があります。
  • graph.lint() の実行
    削除後に graph.lint() を実行し、グラフの整合性に問題がないかチェックします。lint() は、無効なノード参照などを検出するのに役立ちます。
  • より保守的な削除
    もし意図しない削除が頻繁に起こる場合は、eliminate_dead_code() の使用を控え、手動でのグラフ最適化を検討するのも一つの手段です。
• 原因
  グラフの解析が不完全で、実際には後続の演算で必要となるノードが「不要」と判断されてしまうことがあります。これは、複雑な制御フローや動的な挙動を持つモデルで起こりやすいです。

グラフの不整合 (Graph Inconsistency)

• トラブルシューティング
  
  • PyTorch のバージョン確認
    使用している PyTorch のバージョンが最新であるか、または安定版であることを確認します。
  • 再現性の確認
    同じコードで常にエラーが発生するかどうかを確認し、再現可能な場合に PyTorch の Issue Tracker に報告することを検討します。
  • 最小限のコードでの再現
    可能であれば、エラーを再現する最小限のコードを作成し、問題の切り分けを行います。
• 原因
  eliminate_dead_code() の実行中に内部的なエラーが発生し、グラフの構造が壊れてしまうことがあります。これは、PyTorch のバージョン間の互換性問題や、非常に複雑なグラフ構造を持つ場合に稀に起こりえます。

カスタムノードとの互換性 (Compatibility with Custom Nodes)

• トラブルシューティング
  
  • カスタムノードの出力依存性の明示
    カスタムノードの出力が後続のノードで使用されていることを明確にするようにグラフを構築します。
  • replace_all_uses_with() の活用
    もしカスタムノードが削除されて問題がある場合は、削除されたノードの出力を別のノードの出力で置き換える (node.replace_all_uses_with(new_node)) などの手動操作が必要になることがあります。
  • カスタムノードに対する最適化の検討
    カスタムノード自体を FX グラフでより適切に表現する方法を検討します。
• 原因
  torch.fx は、標準的な PyTorch 演算子を追跡しますが、カスタムの演算子 (torch.ops で定義されたものなど) の内部構造までは理解できない場合があります。そのため、カスタムノードが実際には出力に影響を与える場合でも、不要と判断されてしまう可能性があります。

制御フローを含むグラフ (Graphs with Control Flow)

• トラブルシューティング
  
  • 制御フローの単純化
    可能であれば、モデルの構造を単純化し、制御フローの使用を最小限に抑えることを検討します。
  • 動的トレースの検討
    制御フローが複雑な場合は、torch.jit.script や torch.jit.trace などの動的トレース機能の方が適している可能性があります。ただし、これらの機能で生成されたグラフに対して eliminate_dead_code() が常に有効とは限りません。
• 原因
  torch.fx は、現時点では完全に複雑な制御フロー（if 文、ループなど）を静的に解析することは難しい場合があります。そのため、制御フローの内部で定義された変数や演算が、実際には実行される可能性があるにもかかわらず、不要と判断されることがあります。

• PyTorch のドキュメントとコミュニティ
  PyTorch の公式ドキュメントやコミュニティフォーラム（例えば、PyTorch Forums）で同様の問題が報告されていないか検索してみます。
• ログ出力の活用
  グラフの構造やノードの情報を詳しくログ出力するようにコードを修正し、問題発生時の状況を把握しやすくします。
• 段階的な適用
  複雑なグラフに対して一度に eliminate_dead_code() を適用するのではなく、グラフ変換の各段階で適用し、影響を確認することを推奨します。

---

例1：基本的な不要な演算の削除

import torch
import torch.nn as nn
import torch.fx

class SimpleModule(nn.Module):
    def forward(self, x):
        a = x + 1
        b = x * 2  # この結果は使われない
        return a

model = SimpleModule()
graph = torch.fx.symbolic_trace(model)

print("削除前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

graph.eliminate_dead_code()

print("\n削除後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# 削除後のグラフから新しいモジュールを作成 (必要に応じて)
new_module = torch.fx.GraphModule(model, graph)

この例では、SimpleModule の forward メソッド内で b = x * 2 という演算を行っていますが、変数 b は最終的な出力には使われていません。eliminate_dead_code() を実行することで、この不要な乗算演算に関連するノードがグラフから削除されることが期待されます。実行結果を比較すると、削除前のグラフには mul という演算が含まれていますが、削除後のグラフからは消えているはずです。

例2：複数の不要な演算の削除

import torch
import torch.nn as nn
import torch.fx

class MultiDeadCodeModule(nn.Module):
    def forward(self, x):
        y = x + 1
        z = x * 2  # 使われない
        w = y - 3  # 使われない
        return y

model = MultiDeadCodeModule()
graph = torch.fx.symbolic_trace(model)

print("削除前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

graph.eliminate_dead_code()

print("\n削除後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

例3：意図しない削除の可能性 (注意点)

import torch
import torch.nn as nn
import torch.fx

class PotentialIssueModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)
        self.temp_result = None

    def forward(self, x):
        self.temp_result = self.linear(x) # インスタンス変数に保存
        return x + 1

model = PotentialIssueModule()
graph = torch.fx.symbolic_trace(model)

print("削除前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

graph.eliminate_dead_code()

print("\n削除後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# この例では、`self.linear(x)` の結果である `temp_result` は `forward` 関数の戻り値としては使われていません。
# そのため、`eliminate_dead_code()` は `linear` レイヤの呼び出しに関連するノードを削除する可能性があります。
# しかし、`self.temp_result` はインスタンス変数として保存されており、モジュールの状態に影響を与える可能性があります。
# `eliminate_dead_code()` は、グラフのデータフローのみを解析するため、このような副作用までは考慮しないことに注意が必要です。

この例は、eliminate_dead_code() が単純なデータフローの解析に基づいて不要なコードを削除するため、インスタンス変数の更新のような副作用を持つ演算を誤って削除する可能性があることを示唆しています。FX グラフは主に純粋な関数的な演算を対象としており、状態を持つ演算の扱いはより複雑になる場合があります。

例4：グラフ変換後の不要なノードの削除

import torch
import torch.nn as nn
import torch.fx
from torch.fx.passes.fuse_contiguous import fuse_contiguous

class ContiguousModule(nn.Module):
    def forward(self, x):
        y = x.contiguous()
        z = y.view(x.size(0), -1)
        return z

model = ContiguousModule()
graph = torch.fx.symbolic_trace(model)

print("変換前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# contiguous 操作を融合するパスを適用
fused_graph = fuse_contiguous(graph)

print("\n融合後のグラフ:")
for node in fused_graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

# 融合後に不要になった可能性のあるノードを削除
fused_graph.eliminate_dead_code()

print("\n不要コード削除後のグラフ:")
for node in fused_graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs)

new_module = torch.fx.GraphModule(model, fused_graph)

この例では、fuse_contiguous というグラフ変換パスを適用しています。このパスは、.contiguous() の呼び出しを他の演算と融合することがあります。融合後、元々の .contiguous() の演算が不要になる場合があります。eliminate_dead_code() を適用することで、このような変換後に残った不要なノードを削除できます。

---

手動でのノード削除 (Manual Node Removal)

• 欠点
  手間がかかり、グラフの構造を深く理解している必要があります。
• 利点
  不要なノードをより細かく制御できます。自動削除では判断が難しい場合に対応できます。
• 手順
  
  1. グラフのノードをイテレートし (for node in graph.nodes:)、各ノードの属性 (node.op, node.name, node.all_input_nodes, node.users) を確認します。
  2. 削除したいノードを特定します。不要なノードは、その出力 (node.outputs) が他のどのノードの入力としても使われていない (len(node.users) == 0) ことが多いです。ただし、副作用のあるノード（例：インスタンス変数の更新）は、直接の出力依存性がなくても削除すべきでない場合があります。
  3. 特定したノードを graph.erase_node(node) メソッドを使って削除します。
  4. 削除後、graph.lint() を実行してグラフの整合性を確認することを推奨します。

import torch
import torch.nn as nn
import torch.fx

class ManualEliminateModule(nn.Module):
    def forward(self, x):
        a = x + 1
        b = x * 2  # 使われない
        return a

model = ManualEliminateModule()
graph = torch.fx.symbolic_trace(model)

print("削除前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

nodes_to_erase = []
for node in graph.nodes:
    if node.name == "mul":  # 不要な乗算ノードを特定
        nodes_to_erase.append(node)

for node in nodes_to_erase:
    graph.erase_node(node)

graph.lint()

print("\n削除後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

グラフ変換パスの自作 (Custom Graph Transformation Passes)

• 欠点
  パスの作成にはある程度の知識と労力が必要です。
• 利点
  特定の不要なパターンを効率的に削除できます。再利用可能な最適化ロジックを構築できます。
• 手順
  
  1. torch.fx.passes.GraphPass を継承したクラスを作成します。
  2. run_on_graph(self, graph: torch.fx.Graph) メソッドを実装し、グラフをイテレートして不要なパターンを特定し、ノードを削除するロジックを記述します。
  3. 必要に応じて、ノードの置換 (node.replace_all_uses_with()) なども利用できます。
  4. 作成したパスのインスタンスを作成し、グラフに適用します (pass_instance(graph)).

import torch
import torch.nn as nn
import torch.fx
from torch.fx.passes import GraphPass

class RemoveMulPass(GraphPass):
    def run_on_graph(self, graph: torch.fx.Graph) -> torch.fx.Graph:
        nodes_to_erase = []
        for node in graph.nodes:
            if node.op == "mul" and not any(user.op == "output" for user in node.users):
                nodes_to_erase.append(node)

        for node in nodes_to_erase:
            graph.erase_node(node)
        return graph

class MyModule(nn.Module):
    def forward(self, x):
        a = x + 1
        b = x * 2  # 使われない
        return a

model = MyModule()
graph = torch.fx.symbolic_trace(model)

print("変換前のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

remove_mul_pass = RemoveMulPass()
graph = remove_mul_pass(graph)
graph.lint()

print("\n変換後のグラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

• 欠点
  直接的な不要コード削除の制御はできません。
• 利点
  既存の最適化パスを活用することで、効率的な最適化が期待できます。

import torch
import torch.nn as nn
import torch.fx
from torch.fx.passes.fuse_contiguous import fuse_contiguous
from torch.fx.passes.eliminate_dead_code import eliminate_dead_code

class ContiguousDeadModule(nn.Module):
    def forward(self, x):
        y = x.contiguous()
        z = y.view(x.size(0), -1)
        w = x * 2 # 使われない
        return z

model = ContiguousDeadModule()
graph = torch.fx.symbolic_trace(model)

print("初期グラフ:")
for node in graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

fused_graph = fuse_contiguous(graph)
print("\ncontiguous 融合後グラフ:")
for node in fused_graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)

eliminate_dead_code(fused_graph)
print("\n不要コード削除後グラフ:")
for node in fused_graph.nodes:
    print(node.op, node.name, node.all_input_nodes, node.outputs, node.users)