PyTorch

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 
torch.jit.trace は実行時トレースに基づくPython-to-TorchScriptトランスパイラである"Using torch.jit.trace and torch.jit.trace_module, you can turn an existing module or Python function into a TorchScript ScriptFunction or ScriptModule. You must provide example inputs, and we run the function, recording the operations performed on all the tensors." Python. TORCH.JIT.TRACE. pytorch.org. 2023-08-29閲覧.。

torch.jit.trace は実行可能Pythonオブジェクトとその仮入力を引数に取り、これをトレースしながら実行し、トレース結果に基づいてTorchScriptコードを生成する。すなわち実行時トレースに基づいてPythonコードをTorchScriptコードへ変換する。

仮入力を用いたトレースに基づいて生成をおこなうため、利用できるTorchScript機能に制限がある（c.f. torch.jit.script によるコンパイル）。例えば条件分岐は仮入力が満たす片方のパスしか記録されないため、TorchScriptの動的条件分岐として変換されない。すなわちTorchScriptの動的条件分岐は利用できない"Tracing will not record any control-flow like if-statements or loops." Python. TORCH.JIT.TRACE. pytorch.org. 2023-08-29閲覧.。

 FX 
torch.fx はモデルの高レベルグラフ化・変換・コード生成に関するモジュールである"FX is a toolkit for developers to use to transform nn.Module instances." PyTorch. TORCH.FX. 2022-03-23閲覧."torch.fx, a program capture and transformation library for PyTorch" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.。モデルを入力/出力/演算の有向非巡回グラフと見なすIRを定義し"torch.fx represents programs in a DAG-based IR" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.、モデル→IR→モデルの変換とIR編集をサポートする。主に次の3つの機能を提供する"FX consists of three main components: a symbolic tracer, an intermediate representation, and Python code generation." PyTorch. TORCH.FX. 2022-03-23閲覧.。

 Model-to-IR: シンボリックトレースによるPython/PyTorchモデルからのIR生成
 IR-to-IR: IRで表現されたグラフの変換
 IR-to-Model: IRからのPython/PyTorchコード生成

FXはPythonをホスト言語とするIRグラフのメタプログラミングであり"graph-based ... IR ... Program transformations ... is as simple embedded programming languages that are meta-programmed from a host language, predominantly Python" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.、グラフの分析・可視化"torch.fx provides an fx.graph_drawer package, which gives the user the ability to visualize torch.fx graphs with Graphviz" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.、演算子の書き換えによる量子化"Quantization makes use of torch.fx’s graph and GraphModule representation to simultaneously modify the program code and weight values." Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.、グラフ全体を考慮したop fusion"6.2.2 Fusion Optimizations ... torch.fx provides the necessary non-local program context and state modification facilities needed for this transformation with its ahead-of-time, graph-based nature" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.、IRからのアクセラレータコード直接生成" The project was quickly developed using torch.fx’s Python APIs as well as TensorRT’s Python APIs, creating a translation layer between the two." Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.など様々な用途に利用される。

モデルエクスポートのための幅広い演算サポートを目指すTorchScript IR"the purposes of serialization or export. For instance, TorchScript" Reed, et al. (2021). Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python. arxiv.と異なり、FX IRは動的な制御フローを含まず演算ノード単位の操作を前提とした高レベルな表現と、それに付随するシンプルなIR生成・IR操作・モデル生成実装を掲げている。

以下は前向き (forward) 関数の変換の例である：import torch # PyTorchモジュールを読み込み
def f(x, y): return x + y # 前向き関数を作る
fx_traced = torch.fx.symbolic_trace(f) # PyTorchの実行をトレースしてFX Graphを生成
print(fx_traced.graph) # FX GraphのIRを表示
print(fx_traced.code) # FX GraphをPythonコードに再変換して表示

 AOT Autograd 
AOT Autograd (functorch.compile) は前向き関数と後向き関数をコンパイルするための実験的サブシステムであるfunctorch.compile (experimental) — functorch 1.13 documentation PyTorch Foundation。FX GraphをTorchScriptのコンパイラによってコンパイルするための functorch.compile.ts_compile も実験的に提供されているfunctorch.compile.ts_compile — functorch 1.13 documentation PyTorch Foundation。

以下は前向き (forward) 関数の変換の例である：import torch, functorch # PyTorch及びfunctorchモジュールを読み込み
def f(x, y): return x + y # 前向き関数を作る
fx_traced = torch.fx.symbolic_trace(f) # PyTorchの実行をトレースしてFX Graphを生成
ts_compiled = functorch.compile.ts_compile(fx_traced, [torch.Tensor([1,2]), torch.Tensor([3,4])]) # TorchScriptコンパイラでFXグラフをコンパイル
ts_compiled([torch.Tensor([1,2]), torch.Tensor([3,4])]) # コンパイルされた関数を実行

 TorchDynamo 
TorchDynamo (torch._dynamo) は任意のPythonコードをコンパイルのためのサブシステムである。JITコンパイルするためのtorch.compile API、事前コンパイルするためのtorch.export APIもある（後者は#事前コンパイルと量子化参照）。コンパイルできないコードは自動的にPythonインタプリタへとフォールバックされる。2023年にリリースされたPyTorch 2.0で本体へと統合された。

狭義のTorchDynamoはPythonのバイトコードから前述のFX Graphを作るサブシステムであり、FX Graphを各実行環境向けIRへと変換するための各バックエンド（TorchInductor、onnxrtバックエンドなど）と組み合わせてPythonコードのコンパイルを行うことができる。バックエンドを実装しやすくするためのものとして、FX Graphの複合的なAten命令を単純なPrim命令へと置き換えるPrimTorchも存在する。

以下は前向き (forward) 関数の最適化の例である：import torch
def f(x, y): return x + y # 前向き関数を作る
optimized_f = torch.compile(f, backend="inductor") # 前向き関数をTorchDynamo（+ TorchInductor）で最適化
optimized_f(1, 2) # 最適化された前向き関数を実行

また派生プロジェクトとしてバックエンドにOpenXLAのXLAコンパイラを使うためのPyTorch/XLAも開発中となっているが、こちらも試験的にTorchDynamoへと対応しておりTorchDynamo(torch.compile) integration in PyTorch XLA PyTorch FoundationThe Path to Achieve Ultra-Low Inference Latency With LLaMA 65B on PyTorch/XLA PyTorch Foundation、場合によってはTorchDynamo+Inductorよりも高速であるとされる。

 事前コンパイルと量子化 

PyTorch 2.1以降には事前コンパイル (torch.export) が試験的に実装されている。また量子化 (torch.ao.quantization) も試験的にtorch.exportに対応している。

量子化では訓練後量子化（Post Training Quantization; PTQ）と量子化意識訓練 (Quantization Aware Training; QAT) の両方に対応しているQuantization PyTorch Foundation。

 モバイル向け展開 

PyTorchモデルのモバイル向け展開にはPyTorch MobileとExecutorchが存在するExecuTorch PyTorch Foundation。

PyTorch MobileはTorchScriptを使った軽量インタプリタとなっているPyTorch Mobile PyTorch Foundationが、最適化 (torch.utils.mobile_optimizer) torch.utils.mobile_optimizer PyTorch Foundationにも対応している。

一方、Executorchはtorch.compileベースとなっている。

またPyTorchのモデルから別の機械学習フレームワークTensorFlowのモバイル向け実装LiteRT（旧称TensorFlow Lite）のモデルへと変換することも可能であり、これには前述のONNX出力でONNX形式を出力した後にそのファイルをonnx2tfでLiteRTモデルへと変換する方法AI Edge Torch: High Performance Inference of PyTorch Models on Mobile Devices Google 2024年5月14日や、GoogleのAI Edge Torchを使う方法が存在する。

 モデル / ライブラリ 
 Transformer モデル 

性能の高い機械学習モデル「Transformer」が人気となるにつれ、PyTorchでも内部プロジェクト外部プロジェクト合わせ数多くの実装が登場した。その代表的なものには以下が存在する：
 torch.nn.Transformertorch.nn - Transformer Layers PyTorch Foundation - PyTorch 1.2で実装されたNew Releases: PyTorch 1.2, torchtext 0.4, torchaudio 0.3, and torchvision 0.4 PyTorch Foundation 2019年8月8日。
 Accelerated Transformers - 旧Better TransformersPyTorch 2.0: Our next generation release that is faster, more Pythonic and Dynamic as ever PyTorch Foundation 2023年3月15日。HazyResearch社 の FlashAttention や Meta AI社 の xFormers などを取り入れて上記APIを高速化したものA BetterTransformer for Fast Transformer Inference PyTorch Foundation 2022年7月12日であり、PyTorch 2.0で安定版が搭載された。前身のBetter TransformersはHugging FaceのOptimumで簡単に使用することが出来たBetterTransformer - Overview Hugging Face。
 Transformers (Hugging Face) - 様々な応用モデルを実装している。PyTorch以外にも対応している。
 Fairseq (Meta AI) - Transformer派生を含む様々な応用モデルを実装している。またNVIDIAによりこのFairseqに基づくTransformerの実装が「Transformer for PyTorch」として提供されているTransformer for PyTorch　NVIDIA。
 xFormers (Meta AI) - 高速かつ省メモリな複数の実装を含んでいる（CUTLASSを使った実装やOpenAI Tritonを使った実装など）。前述のFairseqや後述のDiffusersでも使えるようになっているFairseq Meta AI動画も音楽もゲームも　広がるAIエンターテインメントの可能性 ITmedia 2022年11月11日。

 拡散モデル 
拡散モデル（確率拡散モデル; Diffusion model）は特に生成AI（ジェネレーティブAI）において人気となり、PyTorchでもその実装が登場した。その代表的なものには以下が存在する：
 Diffusers (Hugging Face) - 単純なU-Netベースの拡散モデルだけでなく、変分オートエンコーダー (VAE) によるを使った潜在拡散モデル (LDM) やその応用の Stable Diffusion (SD)、埋め込み表現にCLIPの代わりにXLM-RobertaベースのAltCLIPを使ってマルチリンガル化した AltDiffusionAltDiffusion Hugging Face、ベクトル量子化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 
torchvisionはPyTorch向けの画像処理向けのライブラリである。画像分類 (Swin Transformer / MaxVit / Swin Transformer V2 / EfficientNet / ConvNeXt / EfficientNetV2 / RegNet / VisionTransformer等)、動画分類 (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) などに対応しているModels and pre-trained weights — Torchvision main documentation PyTorch Foundation。

 torchaudio 
torchaudioはPyTorch向けの音声処理向けのライブラリである。音声認識 (Emformer RNN-T / XLS-R / Conformer RNN-T)、視覚音声認識 (AV-ASR)、テキスト音声合成 (Tacotron2)、音声強調 (MVDR)、ビームフォーミング (DNNBeamformer)、 (Hybrid Demucs)、音声品質測定 (TorchAudio-Squim) 、テキストへの強制位置合わせ (CTC forced alignment) などに対応している。

その他、音声のベクトル化ではwav2vec 2.0、HuBERT、WavLMに対応しているtorchaudio.pipelines PyTorch Foundation。GPUベースのCTCビームサーチ (CUCTCDecoder) にも対応しているNew Library Updates in PyTorch 2.1 PyTorch Foundation 2023年10月4日。

 torchtext 

torchtextはPyTorch向けのテキスト処理ライブラリである。大規模言語モデル (LLM) の XLM-RoBERTaSST-2 Binary text classification with XLM-RoBERTa model PyTorch Foundation や T5T5-Base Model for Summarization, Sentiment Classification, and Translation PyTorch Foundationに対応している。なお大規模言語モデルで文書分類、言語モデリング、質疑応答、機械翻訳などの多様なタスクを行うには、モデルをそれぞれへと適応するためのファインチューニングやプロンプティングが必要となる。

その他、テキストのトークン化では正規表現での置換によるトークン化 (RegexTokenizer)、文字バイト対符号化 (BPE) によるトークン化 (CharBPETokenizer)、バイトレベルBPEによるトークン化（CLIPで使われるCLIPTokenizerやGPT-2で使われるGPT2BPETokenizer）、WordPieceによるトークン化（BERTで使われるBERTTokenizer）、事前学習モデルによるトークン化（T5などに使われるSentencePieceTokenizer）に対応しているtorchtext.transforms PyTorch FoundationSummary of the tokenizers Hugging Face。また単語の埋め込みベクトルへの変換ではGloVe、FastText、CharNGramに対応しているtorchtext.vocab PyTorch Foundation。

2023年9月、開発をメンテナンスモードに移行し、2024年4月開発終了となったtorchtext PyTorch Foundation。

 TorchRL 
TorchRLはPyTorch向けの強化学習ライブラリである。暗黙的Q学習 (IQL)、深層Q学習 (DQN)、 (PPO)、深層決定論的方策勾配法 (DDPG)、双生遅延DDPG (TD3)、Advantage Actor-Critic (A2C)、Soft-Actor-Critic (SAC)、保守的Q学習 (CQL)、ランダム化アンサンブル化ダブルQ学習 (REDQ)、Dreamer、Decision Transformer、人間のフィードバックによる強化学習 (RLHF)に対応しているTorchRL PyTorch Foundation。

 TensorDict 

TensorDictはPyTorch向けの辞書構造ライブラリである。

 TorchRec 

tensordictはPyTorch向けのレコメンデーション（推薦）ライブラリである。

 torchtune 

torchtuneはPyTorch向けの大規模言語モデル (LLM) ファインチューニングライブラリである。

 torchchat / gpt-fast 

torchchatはPyTorch向けのLLMを使ったチャット実装でありIntroducing torchchat: Accelerating Local LLM Inference on Laptop, Desktop and Mobile PyTorch Foundation 2024年7月30日、gpt-fastはPyTorchを使ったLLM実装であるAccelerating Generative AI with PyTorch II: GPT, Fast PyTorch Foundation 2023年11月30日。

 PyTorch Hub 
PyTorch Hub (torch.hub) はGitHubのリポジトリ上に存在するPyTorchの事前学習モデルを簡単に使用できるようにするための仕組みであるtorch.hub — PyTorch 1.13 documentation PyTorch Foundation。リポジトリをこのPyTorch Hubで使えるようにするためにはリポジトリ上に「hubconf.py」を設置する必要がある。

PyTorch Hubに対応しているリポジトリには例えば以下が存在する：
 torchvision ("pytorch/vision")
 Fairseq ("pytorch/fairseq")fairseq · PyPI Facebook
 YOLOv5 ("ultralytics/yolov5")PyTorch Hub Ultralytics

PyTorch Hubでリポジトリ上のモデルを読み込むコードの例：
import torch
yolov5s_model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

なお外部プロジェクトにはGitHub及びPyTorch Hubを使わずに、GitHubに似た独自リポジトリシステムとその読み込みスクリプトを提供するものも存在する（Hugging Face の Hugging Face HubHugging Face Introduces Private Hub Weights & Biases 2022年8月3日Integrate your library with the Hub Hugging Faceなど）。

 TorchMetrics 

Lightning AI製の評価指標（メトリクス）用ライブラリ。

 学習 

 データローダ 
PyTorchに使うことのできるデータローダには以下が存在する：
 DataLoader (torch.utils.data.DataLoader) - PyTorch標準のデータローダ。
 DataLoader2 (torchdata.dataloader2.DataLoader2) - 別パッケージのTorchDataで開発されているDataPipesを基にした実験的なデータローダPyTorch 1.12: TorchArrow, Functional API for Modules and nvFuser, are now available PyTorch Foundation 2022年6月28日。 
 FFCV (Fast Forward Computer Vision; ffcv.loader)
 NVIDIA DALI (Data Loading Library; nvidia.dali.plugin.pytorch) - GPUDirect Storageに対応したデータローダ。PyTorch以外にも対応しているNVIDIA DALI Documentation NVIDIA。実験的なGPUデータ解凍にも対応している (fn.experimental.inflate)。

またデータセットの自動ダウンロードに対応するデータローダとして HuggingFace Datasets Hub のデータセットが使用可能な HuggingFace 製の Datasets も存在する。

 PyTorch Lightning 

PyTorch Lightningは学習から展開までの簡易化のための PyTorch向けサードパーティー製ライブラリである。

 高速化 
より高い効率で学習を行うための様々な標準機能・外部ライブラリが提供されている。以下はその一例である。
+表. PyTorch学習高速化機能・ライブラリ名称手法torch.jit.scriptTorchScript言語での記述 → TorchScriptによる学習の最適化torch.jit.trace仮入力トレースによるTorchScriptへの変換 → TorchScriptによる学習の最適化torch.fx.symbolic_traceFXのTrace機能によるGraphIR化→各バックエンドによる学習の最適化torch._dynamo部分的GraphIR化→PythonJITコンパイラでの部分的バックエンド利用による学習の最適化
それぞれの手法で特徴と適用範囲が異なる。例えばDynamoでは動的条件分岐を扱えるがfx.symbolic_traceでは扱えない"torch.fx.symbolic_trace() ... can’t handle control flow" PyTorch. Why do you need another way of optimizing PyTorch code? - Getting Started - TorchDynamo Overview. pytorch.org. 2023-08-28閲覧.。

 脚注 

 出典 

 参考文献 
 小川雄太郎：「つくりながら学ぶ! PyTorchによる発展ディープラーニング」、マイナビ出版 (2019年7月29日)、ISBN 978-4-8399-7025-3.
 巣籠悠輔：「詳解ディープラーニング 第2版：TensorFlow/Keras・PyTorchによる時系列データ処理」、マイナビ出版、ISBN 978-4-8399-6951-6 (2019年11月27日).

 Eli Stevens、Luca Antiga、Thomas Viehmann：「PyTorch実践入門：デープラーニングの基礎から実装へ」,マイナビ出版(2021年1月30日）、ISBN 978-4-8399-7469-5.
 赤石雅典：「最短コースでわかる PyTorch &深層学習プログラミング」、日経BP（2021年9月17日）、ISBN 978-4-29611032-2.
 Sebastian Raschka、Yuxi (Hayden) Liu、Vahid Mirjalili：「Python 機械学習プログラミング [PyTorch ＆ scikit-learn編]」、インプレス (2022年12月21日)、ISBN 978-4-295-01558-1.

 関連項目 




 Tensor
 人工知能
 ニューラルネットワーク
 ディープラーニング

 外部リンク 

 PyTorchチュートリアル（日本語翻訳版）

Category:応用機械学習
Category:人工ニューラルネットワーク
Category:Pythonライブラリ
Category:オープンソース人工知能
Category:Metaのソフトウェア
Category:Pythonでプログラムされた自由ソフトウェア