为什么越来越多人使用 PyTorch 2？

原创于 2026-08-16 18:50:41 发布 · 161 阅读

2 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

收录于

Python

大家好，我是Java1234_小锋老师。
在这里插入图片描述

如果你最近看过人工智能相关的开源项目，大概率会频繁遇到 PyTorch。它本来就以“写起来像普通 Python”而受到欢迎，到了 PyTorch 2，这种简单直接的体验没有被推翻，运行效率却有了进一步提升。

这也是 PyTorch 2 越来越受欢迎的关键：开发者不必为了速度，把熟悉的代码全部重写一遍。很多时候，只要增加一行 torch.compile()，框架就会尝试在后台优化模型。

PyTorch 2 到底变了什么

PyTorch 2 并不是另起炉灶的新框架。原来会写的张量运算、自动求导、神经网络模块和训练循环，基本都能继续使用。它最受关注的新能力，是以 torch.compile 为入口的编译优化。

简单来说，过去模型更像是“走一步执行一步”；现在框架可以先观察一段计算过程，再把其中适合合并、简化的部分放在一起优化。开发方式仍然灵活，机器执行时却有机会少走弯路。

在这里插入图片描述

大家选择它的几个现实原因

1. 学习成本没有突然增加

PyTorch 2 保留了大家熟悉的编程方式。初学者仍然可以从张量、全连接层和训练循环学起，老项目升级后也不一定要大改。这种“原来的代码还能跑”的兼容性，比单纯增加几个新功能更重要。

2. 加速入口足够简单

使用编译能力最常见的写法只有一行：

model = torch.compile(model)

当然，这不代表所有模型都会立刻变快。第一次运行通常需要编译，模型结构、显卡型号和输入大小也会影响结果。但对开发者来说，尝试优化的门槛确实低了很多。

3. 从实验到落地更顺畅

做研究时，人们希望随时修改网络结构、打印中间结果；做实际项目时，又在意速度和资源消耗。PyTorch 2 尝试把这两种需求放在同一套开发体验里，让代码既方便调整，也有进一步优化的空间。

4. 生态成熟，遇到问题容易找到答案

常见的数据处理、视觉模型、自然语言模型和部署工具，大多已经能够与 PyTorch 配合。教程、示例和社区讨论也很丰富。对个人开发者和团队来说，这意味着试错成本更低。

案例一：用几行代码体验编译加速

下面用一个小型全连接网络演示 torch.compile。为了让计时更可靠，如果使用 CUDA，需要在计时前后等待 GPU 完成任务。

import time
import torch
from torch import nn


class DemoNet(nn.Module):
    """用于演示编译功能的简单全连接网络。"""

    def __init__(self) -> None:
        """初始化网络层。"""
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(1024, 2048),
            nn.ReLU(),
            nn.Linear(2048, 1024),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """执行一次前向计算。"""
        return self.layers(x)


device = "cuda" if torch.cuda.is_available() else "cpu"
model = DemoNet().to(device).eval()
compiled_model = torch.compile(model)
sample = torch.randn(512, 1024, device=device)

# 先预热，避免把首次编译时间混入正式计时
with torch.no_grad():
    for _ in range(5):
        compiled_model(sample)

if device == "cuda":
    torch.cuda.synchronize()

start = time.perf_counter()
with torch.no_grad():
    for _ in range(50):
        compiled_model(sample)

if device == "cuda":
    torch.cuda.synchronize()

print(f"50 次推理耗时：{time.perf_counter() - start:.4f} 秒")

想知道优化是否有效，可以在同一台机器上分别测试 model(sample) 和 compiled_model(sample)。不要只运行一次就下结论，最好先预热，再进行多轮测试。

案例二：搭建一个简单的分类模型

PyTorch 2 的日常训练方式依旧很直观。一个典型过程可以概括为：

在这里插入图片描述

下面的代码使用随机生成的数据演示二分类训练。真实项目中，只需要把随机数据换成自己的数据集即可。

import torch
from torch import nn


class BinaryClassifier(nn.Module):
    """用于二分类任务的小型神经网络。"""

    def __init__(self) -> None:
        """初始化分类网络。"""
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(10, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
        )

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """根据输入特征输出分类分数。"""
        return self.network(features)


torch.manual_seed(42)
features = torch.randn(1000, 10)
labels = (features[:, :3].sum(dim=1) > 0).float().unsqueeze(1)

model = BinaryClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = nn.BCEWithLogitsLoss()

for epoch in range(100):
    optimizer.zero_grad()
    logits = model(features)
    loss = loss_fn(logits, labels)
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 20 == 0:
        print(f"第 {epoch + 1} 轮，损失值：{loss.item():.4f}")