PyTorch
PyTorch 是一个开源 Python 框架,可使用 GPU 张量和自动微分构建并训练神经网络。了解其核心概念和部署流程。
PyTorch 是一个开源框架,用于在 Python 中构建、训练和部署神经网络。它提供由 GPU 加速的张量、自动微分,以及包含层、优化器和损失函数的库;它会逐行即时执行代码,因此模型的编写和调试方式与普通 Python 程序相同。PyTorch 最初由 Facebook AI Research 创建(现为 Meta AI);自 2022 年起,它由 Linux Foundation 旗下的厂商中立组织 PyTorch Foundation 管理,并广泛用于深度学习研究,以及计算机视觉和语言领域的生产模型。
核心概念#
- 张量: 多维数组,类似 NumPy 数组,可以存放在 CPU 或加速器上。PyTorch 通过 CUDA 支持 NVIDIA GPU、通过 ROCm 支持 AMD GPU,并通过 Metal (MPS) 支持 Apple 芯片;在这些设备之间移动数据只需调用一次
.to(device)。 - Autograd: 张量运算执行时,PyTorch 会将其记录到计算图中。调用
.backward()会沿计算图反向遍历,并计算输出相对于每个需要梯度的张量的梯度,这就是训练中的反向传播步骤。 - 即时执行: 每次前向传播都会重新构建计算图,因此模型可以使用 Python 循环、条件语句和打印语句;发生错误时,提示会指向出错的确切代码行。
- 模块和优化器:
torch.nn提供卷积、注意力和归一化等层;torch.optim提供 SGD 和 Adam 等优化器,用计算出的梯度更新权重。 - torch.compile: 它于 PyTorch 2.0 中推出,可将即时执行的程序捕获为计算图并生成融合内核,通常无需更改模型代码,就能加快训练和推理。
下面的示例会在 GPU 可用时运行一个小型计算,并让 autograd 计算其梯度:
import torch
# Use a GPU when PyTorch can see one, otherwise fall back to the CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
# Tensors record the operations applied to them when requires_grad=True
w = torch.tensor([2.0, -1.0], device=device, requires_grad=True)
x = torch.tensor([3.0, 4.0], device=device)
# Ordinary Python runs eagerly: each line executes immediately
loss = (w * x).sum() ** 2
# Autograd walks the recorded graph backward to get d(loss)/dw
loss.backward()
print(f"device={device}, loss={loss.item()}, grad={w.grad.tolist()}")在没有 GPU 的机器上,它会打印 device=cpu, loss=4.0, grad=[12.0, 16.0]。损失为 (2·3 − 1·4)² = 4,autograd 无需手写微积分即可返回其导数 2·(w·x)·x = [12, 16];同样的机制也用于训练具有数十亿个权重的网络。
PyTorch 与其他工具对比#
- TensorFlow: Google 的框架。在 TensorFlow 2 将即时执行设为默认模式之前,它会先构建静态计算图再运行。它仍广泛用于基于 Keras 和 Google 部署工具构建的生产系统。
- JAX: Google 的数值计算库,可通过 XLA 编译的
grad、jit和vmap操作,对纯 Python 函数进行可组合变换;它在 TPU 大规模研究中很受欢迎。 - OpenCV: 一个计算机视觉库,可使用经典算法读取、变换和分析图像。处理流水线通常会使用 OpenCV 加载并预处理帧,然后再交由 PyTorch 模型处理。
- Hugging Face: 一个主要基于 PyTorch 构建的平台和一组库,其 Transformers 库提供预训练语言和视觉模型,并将其封装为 PyTorch 模块。
从研究到部署#
训练好的 PyTorch 模型会保存为检查点,通常是保存权重的 .pt 或 .pth 文件。它可以直接在 Python 服务中运行,但边缘设备和经过优化的服务器通常需要将其转换为针对硬件构建的运行时格式:
Ultralytics YOLO26 使用 PyTorch 编写:yolo26n.pt 是 PyTorch 检查点,训练使用 PyTorch 的 autograd 和优化器运行,而 model.export() 则将模型转换为这些部署格式,包括 ONNX、TensorRT、OpenVINO、Core ML 和 LiteRT。










