PyTorch
PyTorchのautogradは、簡単に計算グラフを定義して勾配を得られるようになっているが、生のautogradは複雑なニューラルネットワークを定義するにはすこし低レベルすぎる場合がある。そのような場合のサポートとしてnnモジュールが提供されている。
PyTorch Lightning は学習から展開までの簡易化のための PyTorch 向けサードパーティー製ライブラリである。
torch.fx はモデルの高レベルグラフ化・変換・コード生成に関するモジュールである。モデルを入力/出力/演算の有向非巡回グラフと見なすIRを定義し、モデル→IR→モデルの変換とIR編集をサポートする。主に次の3つの機能を提供する。
Model-to-IR: シンボリックトレースによるPython/PyTorchモデルからのIR生成
IR-to-IR: IRで表現されたグラフの変換
IR-to-Model: IRからのPython/PyTorchコード生成
FXはPythonをホスト言語とするIRグラフのメタプログラミングであり、グラフの分析・可視化、演算子の書き換えによる量子化、グラフ全体を考慮したop fusion、IRからのアクセラレータコード直接生成など様々な用途に利用される。
モデルエクスポートのための幅広い演算サポートを目指すTorchScript IRと異なり、FX IRは動的な制御フローを含まず演算ノード単位の操作を前提とした高レベルな表現と、それに付随するシンプルなIR生成・IR操作・モデル生成実装を掲げている。
以下は前向き (forward) 関数の変換の例である:
TorchScriptは推論モデル生成のための静的型付け言語である。Pythonのサブセット。TorchScriptプログラムはPythonに依存しない環境(例: C++)で実行可能であり、JITコンパイルにより動作環境での最適化がおこなわれる。JITコンパイラでは演算fusion等がおこなわれる。
なおTorchScriptモデルの保存にはTorchScriptモデル形式が使われている (torch.jit.load / torch.jit.save) が、この形式はPyTorchモデル形式と同じく拡張子に「.pt」が使われている。またTorchScriptを介したONNX形式への変換 (torch.onnx.export) も可能となっている。また外部の LLVM プロジェクトの Torch-MLIR では TorchScript を介してPyTorchモデルのコンパイルを行うことができる。
Lightning AI製の評価指標(メトリクス)用ライブラリ。
性能の高い機械学習モデル「Transformer」が人気となるにつれ、PyTorchでも内部プロジェクト外部プロジェクト合わせ数多くの実装が登場した。その代表的なものには以下が存在する:
torch.nn.Transformer - PyTorch 1.2で実装された。 Accelerated Transformers - 旧Better Transformers。HazyResearch社 の FlashAttention や Meta AI(英語版)社 の xFormers などを取り入れて上記APIを高速化したものであり、PyTorch 2.0で安定版が搭載された。前身のBetter TransformersはHugging FaceのOptimumで簡単に使用することが出来た。
Transformers (Hugging Face) - 様々な応用モデルを実装している。PyTorch以外にも対応している。
Fairseq (Meta AI) - Transformer派生を含む様々な応用モデルを実装している。またNVIDIAによりこのFairseqに基づくTransformerの実装が「Transformer for PyTorch」として提供されている。
xFormers (Meta AI) - 高速かつ省メモリな複数の実装を含んでいる(CUTLASSを使った実装やOpenAI Triton(英語版)を使った実装など)。前述のFairseqや後述のDiffusersでも使えるようになっている。