Facebookの開発したオープンソースのディープラーニングフレームワークであるPyTorchが、バージョン1.8のリリースを発表した。APIがアップデートされ、分散トレーニングが改善された他、AMDのGPUアクセラレータ用のROCmプラットフォームのサポートが追加されている。ドメイン固有ライブラリであるTorchVision、TorchAudio、TorchTextの新バージョンも同時にリリースされた。
PyTorchチームが先日のブログ記事で、リリースのおもな機能について紹介している。新リリースにはROCmプラットフォーム用にビルドされたバイナリが含まれていて、AMD GPUを使用したシステムでのパフォーマンスが向上している。NumPy互換APIの最新版や、パイプライン並列性(pipeline parallelism)や段階的圧縮(gradient compression)といった分散トレーニング用の新機能もある。PythonからPythonへの関数変換を行う新たなツールキットであるtorch.fxのベータ版も提供された。バージョン1.8全体としては、1.7以降、3,000以上のコミットが含まれている。
新しいFXツールキットはJaxとTensorFlowからヒントを得たもので、PyTorch nn.ModuleクラスのサブクラスのPythonコードを対象とする変換メカニズムを開発者に提供する。このツールキットはシンボリックトレーサ、中間表現(intermediate representation)、Pythonコードジェネレータという3つの主要なコンポーネントで構成されていて、これらがModuleサブクラスをGraph表現に変換し、コード内のGraphを修正した上で、新たなGraphをPythonソースコードに変換する。こうして生成されたコードは自動的に、既存のPyTorch Eager-Executionシステムと互換性を持つものになる。PyTorchのドキュメントには、推論目的のモデルの最適化やモデル量子化などを含む、FXのユースケース例がいくつか含まれている。
単一GPUでは処理できないような大規模モデルの問題に対処するため、PyTorchではバージョン1.4からモデル並列トレーニングを導入している。これによって複数のGPU上で、それぞれがモデルパラメータのサブセットをホストする形でのトレーニングが可能になる。しかしながら、このパラダイムの実装では、結果として一度に使用可能なGPUはただひとつであるのが一般的だ。新リリースではGPipeと同じようなパイプライン並列性が導入された。これは、入力されたミニバッチをいくつかのマイクロバッチに分割して複数のGPUにパイプラインすることにより、全体のアイドル時間を短縮(ゼロではない)するものだ。新リリースには分散データ並列フレームワークへの通信フックも取り入れられており、トレーニングの勾配通信ステップの最適化を可能にする。段階的圧縮やPowerSGDなど、ビルド済のフックがいくつか用意されている。
AMD GPUを使用する開発者のために、PyTorchインストーラでは、Radeon Open Compute (ROCm)プラットフォーム用のバイナリビルドを選択するオプションが提供されるようになった。これまではROCmプラットフォーム用のソースコードを構築するか、AMDからバイナリをダウンロードする必要があった。NumPyを使ってスピードアップを図りたいユーザには、高速フーリエ変換(FFT)と一般的な線形代数関数のためのNumPy互換APIが用意されている。さらに、ドメイン固有のPyTorchライブラリもいくつかアップデートされている。TorchVisionにはDetectron2のモバイル版を含む、モバイル向けのアップデートがいくつか加えられた。TorchAudioはI/Oパフォーマンスが向上し、TorchTextのデータセットAPIはPyTorch DataLoaderユーティリティと互換性を持つようになった。
Hacker News上でのリリースに関する議論では、TensorFlowとの比較から、ROCmのサポートではTensorFlowが後れを取ったものの、GoogleのTPUデバイスのサポートでは優位にあることが指摘されている。あるユーザはPyTorchを、次のように称賛している。
私の知る中で、最も印象的なソフトウェアエンジニアリングの成果です ... 信じられないほどの複雑性を、極めてシンプルなインターフェースの背後に隠しています ...
PyTorch 1.8のリリースノートとコードはGitHubで公開されている。