PyTorch
PyTorchは、コンピュータビジョンや自然言語処理で利用されているTorchを元に作られた、Pythonのオープンソースの機械学習ライブラリである。最初はFacebookの人工知能研究グループAI Research lab(FAIR)により開発された。PyTorchはフリーでオープンソースのソフトウェアであり、修正BSDライセンスで公開されている。
さまざまなディープラーニングのソフトウェアがPyTorchを利用して構築されており、その中には、UberのPyro、HuggingFaceのTransformers、Catalystなどがある。
PyTorchは以下の機能を備えている:
強力なGPUサポートを備えた(NumPyのような)テンソル演算
自動微分エンジンの上に構築された深層学習トレーニング (torch.autograd)
自動微分化や自動ベクトル化などの(Google JAXのような)関数変換(functorch / torch.func)
実行時コンパイル (torch.jit)
データ並列処理(#データ並列処理参照)
機械学習モデル形式「ONNX」へのエクスポート(torch.onnx)
FacebookはPyTorchとConvolutional Architecture for Fast Feature Embedding(Caffe2)をメンテナンスしていた。しかし、互換性が無いためPyTorchで定義されたモデルのCaffe2への移行やまたその逆の作業が困難であった。これら2つのフレームワークでモデルを変換することができるように、2017年9月にFacebookとマイクロソフトがOpen Neural Network Exchange(ONNX)プロジェクトを作成した。2018年3月下旬に、Caffe2はPyTorchに併合された。
2019年12月、Preferred Networksは自社開発していたChainerのバージョン7をもって、PyTorchによる研究開発へ順次移行して行くことを発表した
。
PyTorch はテンソルに Tensor (torch.Tensor)と呼ばれるクラスを定義しており、それを均質(homogeneous)な多次元の長方形の数値配列の保存と演算に利用している。PyTorch の Tensor は Numpy の多次元配列 (numpy.ndarray) に似ているが、CUDA が有効な Nvidia のGPU上での演算も可能になっている。
PyTorch には 32bit 浮動小数点数用の FloatTensor や 32bit 整数用の IntTensor など、さまざまな型のTensorサブタイプが存在する。またテンソルにはCPUに配置する torch.*Tensor とGPUに配置する torch.cuda.*Tensor が存在している。それぞれは Tensor.to(...) メソッドなどを用いることで変換することができる。
また PyTorch のテンソルは機械学習の逆伝播に使われる微分のためのパラメータ(勾配データ (Tensor.grad) や 微分関数 (Tensor.grad_fn))を持つことが出来る (requires_grad=Trueの場合、詳しくは後述の#autogradモジュールを参照)。勾配データ (grad) は追加メモリを必要とするため、不要な場合に取り除くことが可能となっている(torch.no_grad() や Tensor.detach() など)。
複素数テンソル(dtype=torch.cfloatのテンソル)や量子化テンソル(dtype=torch.quint8やdtype=torch.qint8などのテンソル)も存在している。量子化テンソルはスケールやゼロポイントなどの量子化パラメータを持っている。
基本のテンソルは長方形であり、非長方形のデータを扱う際には padding で穴埋めする必要があるが、穴の多いデータを穴埋めするのは無駄となるため、非長方形のデータを扱うための疎テンソル (torch.sparse) や ネストされたテンソル (torch.nested) も用意されている。なお、疎テンソル構造はNVIDIA GPUのAmpere以降などでも直接サポートされるようになっているものの、その対応は2:4スパースに限られており、PyTorchでの直接的な実装はまだとなっている。
PyTorch のモデルは基本的に torch.nn.Module の派生クラスとする必要があり、そのtrain() メソッドとeval()メソッドにより、トレーニングモードと評価モードの切り替えが可能となっている。
また to(...) メソッドなどによってモデル全体のパラメータ型の変換が可能であり、それによってモデルパラメータの半精度化(専用のhalf()メソッドもある)やbfloat16浮動小数点数化(専用のbfloat16()メソッドもある)を行うことでメモリ使用量を減らすことが可能となっている。ただしモデルパラメータのより高い圧縮を行うためには torch.quantization などで量子化を行う必要がある。
PyTorch のフロントエンドには Pythonインターフェイスだけでなく C++インターフェイスも存在している。
PyTorchは自動微分と呼ばれるメソッドを利用する。recorderは実行された演算を記録し、その後、勾配の計算を行うときに記録した演算を逆方向にリプレイする。このメソッドは、ニューラルネットワークの構築時に特に強力であるため、順方向のパラメータの微分を計算することで、1エポックの計算にかかる時間を節約することができる。
torch.optimは、ニューラルネットワークの構築時に使用されるさまざまな最適化アルゴリズムを実装したモジュールである。通常使用されるメソッドのほとんどはすでに対応しているため、スクラッチで構築する必要がない。
PyTorchのautogradは、簡単に計算グラフを定義して勾配を得られるようになっているが、生のautogradは複雑なニューラルネットワークを定義するにはすこし低レベルすぎる場合がある。そのような場合のサポートとしてnnモジュールが提供されている。
TorchScriptとして後述。
PyTorch のバックエンドは主に C++ で実装されている (ATen)。ATenは外部から直接使うことも可能となっている (C++ではATenライブラリ、Pythonではtorch.ops.aten)。
また PyTorch にはCPUだけでなくGPUバックエンドもあり、NVIDIA製GPU向けにはCUDAで、AMD製GPU向けにはROCmのHIPで、macOS向けにはMetal Performance Shaders (MPS) で実装されている。CUDAにバージョン毎の互換性の問題があることもあって、PyTorch のバイナリはGPUプラットフォーム毎に別々で提供されている。
また PyTorch では別パッケージとして Torch-TensorRT が用意されており、これを使うことにより PyTorch のバックエンドとして NVIDIA の TensorRT を使うことも可能となっている。
2022年現在はOpenAIのTritonなどをバックエンドとして使用できるTorchDynamoも開発中となっており、これは2023年にPyTorch 2.0に実装される予定となっている。TorchDynamo ではバックエンドの実装を簡略化するために、複合的な Aten 命令を単純な Prims 命令 (torch._prims) へと低下させる機能も提供している。
PyTorchは以下の並列処理に対応している。
マルチスレッドデータ並列処理 (torch.nn.DataParallel) - 2022年現在Pythonのグローバルインタプリタロック (GIL) の問題がありスケールしにくい。
マルチプロセス/マルチマシンデータ並列処理 (DDP; torch.nn.parallel.DistributedDataParallel)
完全共有データ並列処理 (FSDP; torch.distributed.fsdp.FullyShardedDataParallel) - FairScale の FSDP を統合した。
また並列処理のための外部プロジェクトとして Microsoft の DeepSpeed や、それを簡単に使うための HuggingFace の Accelerate も存在する。前述の FairScale 及び DeepSpeed は同じものを実装しており、どちらも最適化状態分割 (OSS; ZeRO Stage 1)、最適化および勾配状態分割 (SDP; ZeRO Stage 2)、完全共有データ並列処理 (FSDP; ZeRO Stage 3) の3つに対応している。
=== データ読み込み ===
外部プロジェクトとしてデータ読み込みを最適化するための FFCV が存在する。
PyTorch で良く使われているファイル形式は Python における Python オブジェクトの標準シリアライズ(直列化)である Pickle形式 (*.pkl) 及びウェイトやバイアスなどのパラメータを無圧縮ZIPアーカイブでまとめたpt形式となっている。PyTorchでのこの形式の読み書きはtorch.load/torch.saveやtorch.jit.load/torch.jit.save(TorchScript向け)などを使うことで可能となっている。この形式はPyTorch内外を含め広く用いられているものの、Pickle形式は仮想マシンの命令コードを記述したものであり様々な操作が実行できるため、セキュリティのリスクが存在している。
そのため PyTorch 標準形式よりも安全で効率的なパラメータ保存形式が開発されてきており、その中の一つ HuggingFace の safetensors 形式は広く普及してきている。
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モデルのコンパイルを行うことができる。
TorchScriptはJITオプティマイザ・JITコンパイラ付きのインタプリタにより実行される。内部ではGraphExecutorによる最適化とinterpreterによる実行がおこなわれる。対応している最適化には、入力型に基づくGraphの特殊化、演算子の融合 (Op fusion)、定数の事前計算、不要ノード削除などがある。
Pythonのコードを直接TorchScriptに変換したり (torch.jit.script)、PyTorchの実行をトレースしてTorchScriptを生成する (torch.jit.trace) ことが可能となっている。
以下は前向き (forward) 関数のスクリプト変換の例である:
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) 関数の変換の例である:
AOT Autograd (functorch.compile) は前向き関数と後向き関数をコンパイルするための実験的サブシステムである。FX GraphをTorchScriptのコンパイラによってコンパイルするための functorch.compile.ts_compile も実験的に提供されている。
以下は前向き (forward) 関数の変換の例である:
TorchDynamo (torchdynamo / torch._dynamo) はPythonコードのコンパイルのためのサブシステムである。コンパイルするためのtorch.compile APIもある。2023年にリリース予定のPyTorch 2.0で本体に統合された。
狭義のTorchDynamoはPythonのバイトコードから前述のFX Graphを作るサブシステムであり、FX Graphを各実行環境向けIRへと変換するための各バックエンド(TorchInductorなど)と組み合わせてPythonコードのコンパイルを行うことができる。バックエンドを実装しやすくするためのものとして、FX Graphの複合的なAten命令を単純なPrim命令へと置き換えるPrimTorchも存在する。
以下は前向き (forward) 関数の最適化の例である:
性能の高い機械学習モデル「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でも使えるようになっている。
拡散モデル(確率拡散モデル; Diffusion model)は特に生成AI(ジェネレーティブAI)において人気となり、PyTorchでもその実装が登場した。その代表的なものには以下が存在する:
Diffusers (Hugging Face) - 単純なU-Netベースの拡散モデルだけでなく、変分オートエンコーダー (VAE) による潜在空間を使った潜在拡散モデル (LDM) やその応用の Stable Diffusion (SD)、埋め込み表現にCLIPの代わりにXLM-RobertaベースのAltCLIPを使ってマルチリンガル化した AltDiffusion、ベクトル量子化VAE (VQ-VAE) による潜在空間を使った VQ-Diffusion、CLIPのテキスト埋め込みを直接使わず画像埋め込みへと変換してからガイドとして使用するDALL-E 2を模したKarlo unCLIP、U-Net部分までTransformerに置き換えたDiT (Diffusion Transformer)、動画生成に対応する Text-to-Video / Text2Video-Zero、オーディオに特化した Audio Diffusion / Latent Audio Diffusion、 MIDIからオーディオを生成する Spectrogram Diffusion、汎用拡散モデルの Versatile Diffusion などを実装している。また追加学習ではTextual Inversion、Text-to-imageファインチューニング、DreamBooth、LoRAなどに対応している。
torchvision は PyTorch 向けの画像処理向けのライブラリである。画像分類 (Swin Transformer / MaxVit / Swin Transformer V2 / EfficientNet / ConvNeXt / EfficientNetV2 / RegNet / VisionTransformer等AlexNet、SqueezeNet、GoogLeNet、MobileNet V2、VGG、MobileNet V3、ShuffleNet V2、MNASNet、DenseNet、Inception V3、ResNet、Wide ResNet、ResNeXtにも対応)、動画分類 (Video ResNet / Video S3D / Video MViT / Video SwinTransformer)、セマンティックセグメンテーション (LRASPP / FCN / DeepLabV3)、物体検出 (SSDlite / SSD / FCOS / RetinaNet / Faster R-CNN)、インスタンスセグメンテーション (Mask R-CNN)、キーポイント検出 (Keypoint R-CNN)、オプティカルフロー (RAFT) などに対応している。
torchaudio は PyTorch 向けの音声処理向けのライブラリである。音声認識 (Emformer RNN-T / XLS-R)、テキスト音声合成 (Tacotron2)、音声強調 (MVDR)、信号分離 (Hybrid Demucs)などに対応している。
その他、音声のベクトル化では wav2vec 2.0、HuBERT、WavLM に対応している。
torchtext は PyTorch 向けのテキスト処理ライブラリである。大規模言語モデル (LLM) の XLM-RoBERTa や T5 に対応している。なお大規模言語モデルで文書分類、言語モデリング、質疑応答、機械翻訳などの多様なタスクを行うには、モデルをそれぞれへと適応するためのファインチューニングやプロンプティングが必要となる。
その他、テキストのトークン化では正規表現での置換によるトークン化 (RegexTokenizer)、文字バイト対符号化 (BPE) によるトークン化 (CharBPETokenizer)、バイトレベルBPEによるトークン化(CLIPで使われるCLIPTokenizerやGPT-2で使われるGPT2BPETokenizer)、WordPieceによるトークン化(BERTで使われるBERTTokenizer)、事前学習モデルによるトークン化(T5などに使われるSentencePieceTokenizer)に対応している。また単語の埋め込みベクトルへの変換では GloVe、FastText、CharNGram に対応している。
torchrl は PyTorch 向けの強化学習ライブラリである。
tensordict は PyTorch 向けの辞書構造ライブラリである。
tensordict は PyTorch 向けのレコメンデーション(推薦)ライブラリである。
PyTorch Hub (torch.hub) は GitHub のリポジトリ上に存在するPyTorchの事前学習モデルを簡単に使用できるようにするための仕組みである。リポジトリをこの PyTorch Hub で使えるようにするためにはリポジトリ上に「hubconf.py」を設置する必要がある。
PyTorch Hub に対応しているリポジトリには例えば以下が存在する:
torchvision ("pytorch/vision")
Fairseq ("pytorch/fairseq")
YOLOv5 ("ultralytics/yolov5")
PyTorch Hubでリポジトリ上のモデルを読み込むコードの例:
なお外部プロジェクトには GitHub 及び PyTorch Hubを使わずに、GitHubに似た独自リポジトリシステムとその読み込みスクリプトを提供するものも存在する(Hugging Face の Hugging Face Hubなど)。
Lightning AI製の評価指標(メトリクス)用ライブラリ。
PyTorchに使うことのできるデータローダには以下が存在する:
DataLoader (torch.utils.data.DataLoader) - PyTorch標準のデータローダ。
DataLoader2 (torchdata.dataloader2.DataLoader2) - 別パッケージのTorchDataで開発されているDataPipesを基にした実験的なデータローダ。
NVIDIA DALI (Data Loading Library; nvidia.dali.plugin.pytorch) - GPUDirect Storageに対応したデータローダ。PyTorch以外にも対応している。実験的なGPUデータ解凍にも対応している (fn.experimental.inflate)。
またデータセットの自動ダウンロードに対応するデータローダとして HuggingFace Datasets Hub のデータセットが使用可能な HuggingFace 製の Datasets も存在する。
Lightning(旧PyTorch Lightning)は学習の簡易化のための PyTorch 向けサードパーティー製ライブラリである。