PyTorch 2.0 来了：我们已知的一切

探索最新的 PyTorch 版本——更快、更 Pythonic，也依然保持高度动态。

已更新 2026年10月6日 · 6分钟阅读

使用 AI 探索

ChatGPT Claude Perplexity

PyTorch 是一个开源、由社区驱动的深度学习框架，能够以灵活高效的方式构建机器学习模型。它拥有友好的用户界面、广泛的社区支持，并与 Python 生态系统无缝集成。

PyTorch 2.0 在核心编译器操作上进行了根本性变革，同时保持了对开发者而言熟悉且易用的体验。此次更新承诺带来更快的性能，以及对动态形状（Dynamic Shapes）和分布式（Distributed）的更广泛支持。

PyTorch 2.0 有哪些新变化？

PyTorch 正在将部分组件从 C++ 迁回 Python，使其更快且更易于修改。在 2.0 版本中，他们引入了 `torch.compile`，这在编译器层面改变了 PyTorch 的运行方式。该功能是可选的，不会影响您现有的代码。

PyTorch 2.0 编译

为为 `torch.compile` 功能提供坚实基础，引入了以下新技术：

TorchDynamo。 一个面向 Python 层的即时（JIT）编译器，专为加速 PyTorch 而设计。通过集成 CPython 的帧评估 API，它在运行时动态修改 Python 字节码，从而加快代码执行。
AOTAutograd。 一套帮助开发者在 PyTorch 上加速模型训练的工具。它会预先跟踪正向和反向图。此外，AOTAutograd 提供了简洁的机制，可使用前沿的深度学习编译器无缝编译提取的计算图。
PrimTorch。 将 PyTorch 的运算符数量从 2000 多个大幅精简到约 250 个原语运算符，显著简化了开发 PyTorch 新特性或后端的过程。
TorchInductor。 一个原生于 PyTorch 的深度学习编译器，可将 PyTorch 模型自动映射为适配多种加速器和后端的生成代码。TorchInductor 使用 OpenAI Triton 作为 GPU 加速的基础构件。

这些新技术均以 Python 编写并支持动态形状。这使新版 PyTorch 运行更快、更灵活、也更易于修改，大幅降低了上手门槛。

代码示例

下面我们快速浏览一下 PyTorch 编译器的简易实现方式。

未使用 torch.compile

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")

未使用 torch.compile

要提升模型性能，只需在模型外包一层 torch.compile，即可获得编译后的模型。即插即用。

import torch
model = torch.hub.load("pytorch/vision", "resnet50", weights="IMAGENET1K_V2")
compiled_model = torch.compile(model)

通过参加 Deep Learning with PyTorch 课程，学习如何使用 PyTorch 库构建深度学习模型。

您可以在不做任何改动的情况下直接训练模型。

import torch
model = torch.compile(model)
for batch in dataloader:
    run_epoch(model, batch)

或者您也可以进行推理。

model = torch.compile(model)
model(**input)

torch.compile() 还提供了附加参数：

mode：您可以指定编译器在编译时应优化的目标。
dynamic：用于启用动态形状的代码路径。
fullgraph：将程序编译为单一计算图。
backend：默认使用 TorchInductor，但您也可以指定其他可用的编译后端。

def torch.compile(model: Callable,
  *,
  mode: Optional[str] = "default",
  dynamic: bool = False,
  fullgraph:bool = False,
  backend: Union[str, Callable] = "inductor",
  **kwargs
) -> torch._dynamo.NNOptimizedModule

通过我们的 PyTorch 教程学习 PyTorch API 基础，并从零构建一个简单的神经网络。

基准测试

在上手 PyTorch 2.0 时，开发者使用了 163 个开源（46 个 HuggingFace Transformers、61 个 TIMM 和 56 个 TorchBench）模型来为新编译功能创建性能基准。该基准涵盖了图像分类、图像生成、语言建模、推荐系统与强化学习等任务。

结果显示，在 NVIDIA A100 GPU 上训练时，性能有显著提升。

注意：当前默认后端仅支持 CPU 以及 Nvidia Volta 和 Ampere 系列 GPU。

NVIDIA A100 GPU 上的 PyTorch 编译器基准测试

这只是开始，您将在后续更新中看到更多性能和可扩展性方面的改进。

如何安装 PyTorch 2.0

您可以通过 pip 轻松安装较新的 PyTorch 版本。

将以下命令复制粘贴到您的终端中。

适用于 GPU：CUDA 11.8

事实证明，较新的 GPU 版本表现出了显著更好的性能。

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118

适用于 GPU：CUDA 11.7

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117

适用于 CPU：

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu

用于验证：

git clone https://github.com/pytorch/pytorch
cd tools/dynamo
python verify_dynamo.py

用 PyTorch 2.0 加速 Hugging Face

让我们尝试使用 torch.compile 来加速 Hugging Face transformers。只需一行装饰器，您就能让 Hugging Face 代码运行得更快。

注意：使用 torch.compile()，我们在训练中观察到 30%-200% 的性能提升——参见 TorchDynamo 性能看板。

在这个示例中，我们将在 “dolly-v2-3b” 大语言模型上应用 torch.compile 以获得更快的推理。要在 Google Colab 中运行代码，我们首先需要安装必要的 Python 库。

%%capture
%pip install transformers accelerate xformers

接着，我们将使用 Hugging Face transformers 下载并加载分词器和语言模型。之后，我们会将 nn.Module 传给 torch.compile() 函数。

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationConfig,
    pipeline,
)


tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b")
model = AutoModelForCausalLM.from_pretrained(
    "databricks/dolly-v2-3b", device_map="auto", torch_dtype=torch.bfloat16
)
model = torch.compile(model) #only line of code is required

最后一步，我们将使用 tokenizer 将文本转换为 token，传入 model.generate，然后用 tokenizer.batch_decode 将生成的输出解码为文本。

prompt = "I love you because"
inputs = tokenizer(prompt, return_tensors="pt").to(device="cuda:0")


# Generate
generate_ids = model.generate(inputs.input_ids, max_length=50)
tokenizer.batch_decode(
    generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
)[0]

可以看到，“dolly-v2” 通过补充 “you are a good person…” 完成了句子。

"I love you because you are a good person.
You are kind, you help others, you are honest, you are loyal, you are humble, you are humble, you are humble.
You are a good person."

它同样适用于 Hugging Face 的 pipeline。只需提供任务类型、模型和分词器。

generator = pipeline("text-generation", model= model,tokenizer=tokenizer)
generator("What is the name of Germany's Capital?")

输出：

[{'generated_text': "What is the name of Germany's Capital? 
The name of Germany's Capital is Berlin."}]

该编译功能可与 transformers、accelerate 和 TIMM 等 Python 库开箱即用。

如果您还不确定从何处开始深度学习与 AI 职业之路，Machine Learning Scientist with Python 职业路径是极佳的起点。它全面概述了助您胜任机器学习科学家所需的关键 Python 技能。

Author

Abid Ali Awan

作为一名持证的数据科学家，我热衷于利用前沿技术打造创新的机器学习应用。凭借在语音识别、数据分析与报告、MLOps、对话式人工智能以及自然语言处理方面的扎实背景，我不断打磨构建智能系统的能力，力求带来切实影响。除技术专长外，我也擅长沟通，能够将复杂概念提炼为清晰、简明的表述。因此，我成为数据科学领域备受关注的博主，与不断壮大的数据专业人士社区分享洞见与实践经验。目前，我专注于内容创作与编辑，借助大语言模型打造有力且吸引人的内容，帮助企业与个人更好地发挥数据价值。

主题

Python