【PyTorch】TorchScript変換時のエラー解決法

はじめに

PyTorchで発生するTorchScript変換エラーの原因と解決法を解説します。

症状

PyTorchモデルをTorchScript(torch.jit.script/trace)で変換する際に、型エラーやグラフ構築エラーが発生します。

原因

1. 動的制御フローの非対応

if文やforループがテンソルの値に依存する場合、traceでは捕捉できません。

2. 型アノテーションの不足

TorchScriptは静的型付けのため、型情報が必要です。

解決方法

手順1: torch.jit.traceの使用

# trace: 入力例を使って計算グラフを記録
example_input = torch.randn(1, 3, 224, 224)
traced = torch.jit.trace(model, example_input)
traced.save('model_traced.pt')

手順2: torch.jit.scriptの使用

# script: コードを静的に解析
@torch.jit.script
def process(x: torch.Tensor) -> torch.Tensor:
    if x.sum() > 0:
        return x * 2
    return x

まとめ

PyTorchのTorchScript変換エラーについて解説しました。この記事が問題解決の参考になれば幸いです。