PyTorchは、コンピュータビジョンや自然言語処理で利用されているを元に作られた、Pythonのオープンソースの機械学習ライブラリである。最初はFacebookの人工知能研究グループAI Research lab（FAIR）により開発された。PyTorchはフリーでオープンソースのソフトウェアであり、修正BSDライセンスで公開されている。
さまざまなディープラーニングのソフトウェアがPyTorchを利用して構築されており、その中には、UberのPyro、HuggingFaceのTransformers、Catalystなどがある。
PyTorchは以下の機能を備えている： 強力なGPUサポートを備えた（NumPyのような）テンソル演算 自動微分エンジンの上に構築された深層学習トレーニング (torch.autograd) 自動微分化や自動ベクトル化などの（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のテンソル PyTorchはテンソルにTensor (torch.Tensor) と呼ばれるクラスを定義しており、それを均質（）な多次元の長方形の数値配列の保存と演算に利用している。PyTorchのTensorはNumPyの多次元配列 (numpy.ndarray) に似ているが、CUDAが有効なNvidiaのGPU上での演算も可能になっている。NumPyの配列からPyTorchのテンソルへと変換するための専用APIも存在する (torch.from_numpy)。なおPyTorch 2.1以降はPyTorchだけでなくNumPyを使ったコードのGPU向けコンパイルにも対応している。
PyTorchには32ビット浮動小数点数用の FloatTensor 16ビット浮動小数点数用のHalfTensor 、32ビット整数用の IntTensor など、さまざまな型のTensorサブタイプが存在するtorch.tensor PyTorch Foundation。またテンソルにはCPUに配置する torch.*Tensor とGPUに配置する torch.cuda.*Tensor が存在している。それぞれは Tensor.to(...) メソッドなどを用いることで変換することができる。
また PyTorch のテンソルは機械学習の逆伝播に使われる微分のためのパラメータ（勾配データ (Tensor.grad) や微分関数 (Tensor.grad_fn)）を持つことが出来る（requires_grad=Trueの場合、詳しくは後述の#autogradモジュールを参照）Autograd PyTorch Foundation。勾配データ (grad) は追加メモリを必要とするため、不要な場合に取り除くことが可能となっている（torch.no_grad() や Tensor.detach() など）。
複素数テンソル（dtype=torch.cfloatのテンソル）Complex Numbers PyTorch Foundationや量子化テンソル（dtype=torch.quint8やdtype=torch.qint8などのテンソル）Quantization PyTorch Foundationも存在している。量子化テンソルはスケールやゼロポイントなどの量子化パラメータを持っている。なお4ビット量子化向けのdtype=torch.quint4x2のテンソルもあるが、2023年現在その対応は一部（EmbeddingBag命令）のみとなっている。また8ビット浮動小数点数向けの torch.float8_e4m3fn 及び torch.float8_e5m2 のテンソルにも対応しているが対応命令はごく一部となっている[（8ビット浮動小数点数はNVIDIA H100などの一部GPUのTransformer Engineがネイティブに対応しているUsing FP8 with Transformer Engine NVIDIA）。
基本のテンソルは長方形であり、非長方形のデータを扱う際にはpaddingで穴埋めする必要があるが、穴の多いデータを穴埋めするのは無駄となるため、非長方形のデータを扱うための疎テンソル (torch.sparse) やネストされたテンソル (torch.nested) も用意されている。なお、疎テンソル構造はNVIDIA GPUのAmpere以降などでも直接サポートされるようになっているものの、その対応は2:4スパース（半構造化スパース）に限られているNVIDIA Ampereにおけるプルーニング対応の特徴 Impress 2020年7月20日が、PyTorch 2.1以降はその疎半構造化テンソルにも試験的に対応している (SparseSemiStructuredTensor)Sparse Semi-Structured Tensors PyTorch FoundationPyTorch 2.1: automatic dynamic shape compilation, distributed checkpointing。
PyTorchのモデル
PyTorchのモデルは基本的にtorch.nn.Moduleの派生クラスとする必要があり、そのtrain() メソッドとeval()メソッドにより、トレーニングモードと評価モードの切り替えが可能となっている。
またto(...)メソッドなどによってモデル全体のパラメータ型の変換が可能でありModule PyTorch Foundation、それによってモデルパラメータの半精度化（専用のhalf()メソッドもある）やbfloat16化（専用のbfloat16()メソッドもある）を行うことでメモリ使用量を減らすことが可能となっている。ただしモデルパラメータのより高い圧縮を行うためには torch.quantization などで量子化を行う必要がある。
フロントエンド PyTorchのフロントエンドには PythonインタフェースだけでなくC++インタフェースも存在している。
autogradモジュール PyTorchは自動微分と呼ばれるメソッドを利用する。recorderは実行された演算を記録し、その後、勾配の計算を行うときに記録した演算を逆方向にリプレイする。このメソッドは、ニューラルネットワークの構築時に特に強力であるため、順方向のパラメータの微分を計算することで、1エポックの計算にかかる時間を節約することができる。
optimモジュール torch.optimは、ニューラルネットワークの構築時に使用されるさまざまな最適化アルゴリズムを実装したモジュールである。通常使用されるメソッドのほとんどはすでに対応しているため、スクラッチで構築する必要がない。
nnモジュール PyTorchのautogradは、簡単に計算グラフを定義して勾配を得られるようになっているが、生のautogradは複雑なニューラルネットワークを定義するにはすこし低レベルすぎる場合がある。そのような場合のサポートとしてnnモジュールが提供されている。
jitモジュール TorchScriptとして後述。
バックエンド
PyTorchのバックエンドは主にC++で実装されている (ATen)。ATenは外部から直接使うことも可能となっている (C++ではATenライブラリTensor Basics PyTorch Foundation、Pythonではtorch.ops.aten)。
またPyTorchにはCPUだけでなくGPUバックエンドもあり、NVIDIA製GPU向けにはCUDAで、AMD製GPU向けにはのHIPで、インテル製GPU向けにはSYCLでGetting Started on Intel GPU PyTorch FoundationPyTorch 2.5 Released With Improved Intel GPU Support Phoronix 2024年10月17日、macOS向けにはMetal Performance Shaders (MPS) で実装されている。CUDAにバージョン毎の互換性の問題があることもあって、PyTorch のバイナリはGPUプラットフォーム毎に別々で提供されている。
AI専用チップ向けのバックエンドでは外部プロジェクトとしてIntel Gaudi向けのhabana_frameworks.torchIntel Gaudi PyTorch Python API (habana_frameworks.torch) — Gaudi Documentation 1.18.0 documentation Intel、Huawei Ascend 向けのtorch_npuなどが存在するFacilitating New Backend Integration by PrivateUse1 PyTorch Foundation。PyTorchではこれら外部バックエンドを登録するための仕組み「PrivateUse1」が用意されているChallenges and Efforts in PyTorch Multi-Device Integration: Compatibility, Portability, and Integration Efficiencies PyTorch Foundation 2024年9月18日。
PyTorch 2.0以降にはOpenAIのTritonなどをバックエンドとして使用できるTorchDynamoも統合されている（#TorchDynamoを参照）。TorchDynamoではバックエンドの実装を簡略化するために、複合的な Aten 命令を単純な Prims 命令 (torch._prims) へと低下させる機能も提供しているPyTorch 2.0 PyTorch Foundation。
またPyTorchでは別パッケージとしてTorch-TensorRTが用意されており、これを使うことにより PyTorchのバックエンドとしてNVIDIAのTensorRTを使うことも可能となっているInstallation PyTorch FoundationTorch-TensorRTでPyTorchの推論を最大 6 倍高速化 NVIDIA 2021年12月2日。
並列処理 PyTorchは標準で並列処理機能を内蔵しているものの一部のみとなっている。外部プロジェクトのはより高度な並列処理を実装している。並列処理を簡単に使うための補助ライブラリとしてHuggingFace AccelerateやLightning Fabric（旧Lightning Lite）が存在する。
また既存モデルでは推論の並列処理を簡単に使うためのDeepSpeed も存在するDeepSpeed-MII: instant speedup on 24,000+ open-source DL models with up to 40x cheaper inference Microsoft 2022年10月10日DeepSpeed Deep Dive — Model Implementations for Inference (MII) Towards Data Science 2022年11月17日。
またコンピュータ・クラスターを使った学習向けでは外部プロジェクトにMosaicMLのComposerが存在し、このComposerはクラウドコンピューティングを使った学習とも相性が良いとされるReduce deep learning training time and cost with MosaicML Composer on AWS Amazon Web Services 2022年10月24日Cloudflare R2 and MosaicML enable training LLMs on any compute, anywhere in the world, with zero switching costs Cloudflare 2023年5月16日Train and deploy generative AI faster with MosaicML and Oracle オラクル 2023年5月16日。
データ並列 (DP) PyTorchは以下のデータ並列処理に対応している: マルチスレッドデータ並列処理 (torch.nn.DataParallel) - 2024年現在Pythonのグローバルインタプリタロック (GIL) の問題がありスケールしにくいGetting Started with Distributed Data Parallel — PyTorch Tutorials 2.3.0+cu121 documentation PyTorch Foundation。 マルチプロセス/マルチマシンデータ並列処理 (DDP; torch.nn.parallel.DistributedDataParallel) 完全共有データ並列処理 (FSDP; torch.distributed.fsdp.FullyShardedDataParallel) - FairScaleのFSDPを統合したIntroducing PyTorch Fully Sharded Data Parallel (FSDP) API PyTorch Foundation 2022年3月14日。
前述のFairScale及びDeepSpeedは、どちらも最適化状態分割 (OSS; ZeRO Stage 1)、最適化および勾配状態分割 (SDP; ZeRO Stage 2)、完全共有データ並列処理 (FSDP; ZeRO Stage 3) の3つに対応しているZeRO MicrosoftEfficient Memory management - FairScale documentation Facebook Research。加えてDeepSpeedはZeRO Stage 3よりもノード間通信効率の良いZeRO++にも対応しているZeRO++ Microsoft。
その他DeepSpeedにはデータをCPUメモリやNVMeなどへとオフロードするための ZeRO-InfinityZeRO-Infinity and DeepSpeed: Unlocking unprecedented model scale for deep learning training Microsoft（及びそのサブセットでホストCPUメモリのみ対応の旧来の ZeRO-OffloadZeRO-Offload Microsoft）も搭載されている。
パイプライン並列 (PP) PyTorchは以下のパイプライン並列処理に対応している: パイプライン並列処理 (torch.distributed.pipeline) - nn.Sequentialにのみ対応。FairScaleの実装 (Fairscale.nn.Pipe) を統合したもので、元々は torchgpipe の GPipe実装に由来するPipeline Parallelism PyTorch FoundationPipeline Parallelism - FairScale documentation Facebook Research。
PyTorchモデルの自動パイプライン並列処理に向けては PiPPy が開発中となっているPiPPy PyTorch Foundation。またDeepSpeedもパイプライン並列処理に対応している (deepspeed.pipe)Pipeline Parallelism - DeepSpeed Microsoft。
また推論のパイプライン並列処理にはPetalsPETALS: Collaborative Inference and Fine-tuning of Large Models p.2 Alexander Borzunov et al. 2023年やFlexGenFlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU p.6 Ying Sheng et al. 2023年も存在する。
テンソル並列 (TP) PyTorchは以下のテンソル並列処理に対応している: テンソル並列処理 (torch.distributed.tensor.parallel)Tensor Parallelism - torch.distributed.tensor.parallel PyTorch Foundation
またDeepSpeedもDeepSpeed-Inferenceで推論におけるテンソル並列処理に対応しているModel Parallelism — transformers 4.11.3 documentation Hugging Face。
エキスパート並列 エキスパート並列は複数エキスパート（専門家）の混合 () レイヤー（Switch Transformerなど）を用いたモデルの並列処理を行うMixture of Experts - DeepSpeed Microsoft。
このエキスパート並列には DeepSpeed MoE が対応している。
推論の自動ウェイトオフロード
HuggingFaceのAccelerateは大規模モデルの推論を処理するためにウェイトの自動オフロードに対応しているHandling big models for inference HuggingFaceHow 🤗 Accelerate runs very large models thanks to PyTorch HuggingFace 2022年9月27日。
またDeepSpeed も推論でのデータオフロードにも対応している（DeepSpeed ZeRO-Inference）ZeRO-Inference: Democratizing massive model inference Microsoft 2022年9月9日。
その他 PyTorchはCUDAの複数ストリームを使った明示的なGPU並列処理にも対応しているDocs > CUDA semantics - CUDA streams PyTorch Foundation。
ファイル形式
PyTorchで良く使われているファイル形式はPythonにおけるPythonオブジェクトの標準シリアライズ（直列化）であるPickle形式 (*.pkl) 及びウェイトやバイアスなどのパラメータを無圧縮ZIPアーカイブでまとめたpt形式となっている。PyTorchでのこの形式の読み書きはtorch.load/torch.saveやtorch.jit.load/torch.jit.save（TorchScript向け）などを使うことで可能となっている。この形式はPyTorch内外を含め広く用いられているものの、Pickle形式は仮想マシンの命令コードを記述したものでありPython pickle serialization format: format specification Kaitai Project様々な操作が実行できるため、セキュリティのリスクが存在しているPickle Scanning HuggingFace。
そのためPyTorch標準形式よりも安全で効率的なパラメータ保存形式が開発されてきており、その中の一つHuggingFaceのsafetensors形式は広く普及してきている。
ONNX出力
ONNX出力 (torch.onnx) には従来のTorchScript経由によるもの (torch.onnx.export) と新しいTorchDynamo経由によるもの (torch.onnx.dynamo_export) があるtorch.onnx PyTorch Foundation。前者は#TorchScriptを参照。
TorchScript TorchScriptは推論モデル生成のための静的型付き言語である"TorchScript is a way to create serializable and optimizable models from PyTorch code." PyTorch. TorchScript.。Pythonのサブセット"TorchScript is a statically typed subset of Python" PyTorch. TorchScript Language Reference.。TorchScriptプログラムはPythonに依存しない環境（例: C++）で実行可能であり"TorchScript program that can be run independently from Python, such as in a standalone C++ program." PyTorch. TorchScript.、JITコンパイルにより動作環境での最適化がおこなわれる。JITコンパイラでは演算fusion等がおこなわれる"In PyTorch 1.10, we’ve added an LLVM-based JIT compiler for CPUs that can fuse together sequences of torch library calls" PyTorch. PyTorch 1.10 Release.。
なおTorchScriptモデルの保存にはTorchScriptモデル形式が使われている (torch.jit.load / torch.jit.save) が、この形式はPyTorchモデル形式と同じく拡張子に「.pt」が使われているSaving and Loading Models — PyTorch Tutorials 1.12.1+cu102 documentation PyTorch Foundation。またTorchScriptを介したONNX形式への変換 (torch.onnx.export) も可能となっているtorch.onnx — PyTorch 1.13 documentation PyTorch Foundation。また外部のLLVMプロジェクトのTorch-MLIRではTorchScriptを介してPyTorchモデルのコンパイルを行うことができるMy spring internship – torch-mlir eager mode, OPT and blowing away the main git repo nod.ai 2022年6月20日。
実行 TorchScriptはJITオプティマイザ・JITコンパイラ付きのインタプリタにより実行される"TorchScript is executed using an interpreter attached to a JIT-optimizer and compiler." PyTorch. JIT Technical Overview. 2022-03-19閲覧.。内部ではGraphExecutorによる最適化とinterpreterによる実行がおこなわれる。対応している最適化には、入力型に基づくGraphの特殊化"The executor specializes the Graph for the particular set of inputs." PyTorch. JIT Technical Overview. 2022-03-19閲覧.、演算子の融合 (Op fusion)、定数の事前計算"It propagates constants, pre-computing as much as possible" PyTorch. JIT Technical Overview. 2022-03-19閲覧.、不要ノード削除"Eliminating dead code" PyTorch. JIT Technical Overview. 2022-03-19閲覧.などがある。
生成 TorchScriptコードは直接記述 (torch.jit.script) およびPyTorch実行トレース (torch.jit.trace) により生成される。
以下は前向き (forward) 関数のスクリプト変換の例である：import torch # PyTorchモジュールを読み込み def f(x, y): return x + y # 前向き関数を作る ts_converted = torch.jit.script(f) # Pythonコードから直接TorchScriptに変換 print("Converted:\n%s" % ts_converted.code) # 変換したTorchScriptを表示 ts_traced = torch.jit.trace(f, (torch.tensor([1,2]), torch.tensor([3,4]))) # PyTorchの実行をトレースしてTorchScriptを生成 print("Traced:\n%s" % ts_traced.code) # 生成したTorchScriptを表示
torch.jit.save(ts_converted, "converted.pt") # 変換したTorchScriptをTorchScriptモデル形式で保存 torch.onnx.export(ts_converted, (torch.tensor([1,2]), torch.tensor([3,4])), "converted.onnx") # 上記をONNX形式で出力
torch.jit.script torch.jit.script はTorchScriptコンパイラのPythonラッパー関数である"Scripting a function or nn.Module will inspect the source code, compile it as TorchScript code using the TorchScript compiler, and return a ScriptModule or ScriptFunction." PyTorch. TORCH.JIT.SCRIPT. pytorch.org. 2023-08-29閲覧.。
torch.jit.script はPythonオブジェクトとして書かれたTorchScriptコードをコンパイラへ渡し、コンパイル結果へPythonからアクセスするためのラッパーオブジェクトを返す。Pythonデコレータとしても利用できる。
TorchScriptコードのコンパイルをおこなうため、全てのTorchScirpt機能を利用できる（c.f. torch.jit.trace によるトレース）。例えば動的条件分岐を記述できる。
torch.jit.trace