PyTorch 是由 Meta（原 Facebook）AI 研究团队于 2016 年开源的深度学习框架，以其动态计算图（Define-by-Run）机制、Pythonic 编程范式和强大的 GPU 加速能力，迅速成为学术界和工业界最受欢迎的深度学习框架之一。
| 时间 | 里程碑 | 说明 |
|---|---|---|
| 2016年12月 | PyTorch 0.1 发布 | 基于 Torch 框架，使用 Python 前端替代 Lua |
| 2018年4月 | PyTorch 0.4 发布 | 合并 Variable 和 Tensor，简化 API |
| 2018年12月 | PyTorch 1.0 发布 | 引入 JIT 编译器、TorchScript 和 C++ API |
| 2019年9月 | PyTorch 1.3 发布 | 支持移动端部署、量化、TPU |
| 2020年5月 | PyTorch 1.5 发布 | 引入 torch.fx 程序化变换 |
| 2021年10月 | PyTorch 1.10 发布 | CUDA Graphs 支持，增强分布式训练 |
| 2022年3月 | PyTorch 1.11 发布 | torch.cuda.amp 原生支持混合精度 |
| 2022年9月 | PyTorch 1.13 发布 | torch.compile 预览版，引入 Dynamo |
| 2023年3月 | PyTorch 2.0 正式发布 | torch.compile 稳定版，训练提速 40%+ |
| 2024年 | PyTorch 2.x 系列 | 持续优化编译能力、分布式训练、模型量化 |
| 对比维度 | PyTorch | TensorFlow |
|---|---|---|
| 计算图机制 | 动态图（Define-by-Run） | 静态图（定义为 2.x 支持 Eager） |
| 调试体验 | Python 原生调试器 | 需要 tfdbg 专用调试器 |
| API 简洁性 | Pythonic，设计哲学简洁 | API 层次复杂，版本间变化大 |
| 学术界采用率 | 92%（2024 NLP/DL 顶会论文） | 约 5% |
| 工业部署 | TorchServe、ONNX、TorchScript | TF Serving、TFLite 成熟 |
| 分布式训练 | DDP/FSDP（原生） | Distribution Strategy |
| 移动端支持 | Torch Mobile（有限） | TFLite（较成熟） |
| 可视化 | TensorBoard 集成 | TensorBoard 原生 |
| 社区生态 | 灵活、成长快 | 成熟、企业级支持 |
数据来源：2024年各大顶会（NeurIPS、ICML、ICLR）论文框架使用统计。
Tensor 是 PyTorch 中最基本的数据结构，类似于 NumPy 的 ndarray，但支持 GPU 加速与自动微分。
import torch
# 创建不同维度的张量
scalar = torch.tensor(5.0) # 0 维：标量
vector = torch.tensor([1, 2, 3]) # 1 维：向量
matrix = torch.tensor([[1, 2], # 2 维：矩阵
[3, 4]])
tensor3d = torch.zeros(2, 3, 4) # 3 维：形状 (2,3,4)
print(f"标量: {scalar}, shape: {scalar.shape}")
print(f"向量: {vector}, shape: {vector.shape}")
print(f"矩阵: {matrix}, shape: {matrix.shape}")
print(f"3D 张量: shape: {tensor3d.shape}")
核心属性：
| 属性 | 含义 | 示例 |
|---|---|---|
shape |
张量形状 | (64, 3, 224, 224) 表示 64 张 RGB 图片 |
dtype |
数据类型 | torch.float32, torch.int64 |
device |
存储设备 | 'cpu', 'cuda:0' |
requires_grad |
是否需要梯度 | 训练参数设为 True |
PyTorch 的自动微分机制是其核心引擎。它通过记录所有张量操作来构建计算图，然后自动计算梯度。
# 计算梯度示例：y = (x + 2)^2
x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x + 2) ** 2
# 标量输出才能反向传播
z = y.sum()
z.backward() # 计算梯度
print(f"x: {x}")
print(f"y = (x+2)^2: {y}")
print(f"dz/dx = 2(x+2): {x.grad}")
# 输出: dz/dx = 2(x+2): [6., 8., 10.]
梯度计算验证：
对于 ，梯度为 。
当 时：
与传统 TensorFlow 1.x 先定义静态图再执行的模式不同，PyTorch 每次前向传播都重新构建计算图，这使得：
if/for 等 Python 语句自然融入def dynamic_net(x, use_dropout=True):
"""演示动态计算图的灵活性"""
h = torch.relu(x @ w1 + b1)
if use_dropout and training: # Python 控制流
h = torch.dropout(h, p=0.5)
for _ in range(num_layers): # 循环
h = torch.relu(h @ w2 + b2)
return h @ w3 + b3
静态图 vs 动态图对比：
| 特性 | 静态图（TensorFlow 1.x） | 动态图（PyTorch） |
|---|---|---|
| 构建方式 | 先声明再执行 | 即声明即执行 |
| 控制流 | 需要特殊 API（tf.cond） |
Python 原生 |
| 优化机会 | 编译时可全局优化 | 运行时优化（TorchDynamo 弥合差距） |
| 调试 | 难（图级抽象） | 易（Python 原生） |
| 序列长度可变 | 难（需要 padding） | 自然支持 |
| 启动延迟 | 低（图已编译） | 高（需要 JIT 最新版改善） |
nn.Module 基础所有神经网络层都继承自 torch.nn.Module：
import torch.nn as nn
import torch.nn.functional as F
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 输入通道1，输出32
self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
self.fc1 = nn.Linear(64 * 5 * 5, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = x.view(x.size(0), -1) # flatten
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
model = SimpleCNN()
print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")
| 层类型 | 类名 | 输入形状 | 输出形状 | 参数 |
|---|---|---|---|---|
| 线性层 | nn.Linear(m, n) |
(batch, m) |
(batch, n) |
权重: , 偏置: |
| 卷积层 | nn.Conv2d(c_in, c_out, k) |
(B, c_in, H, W) |
(B, c_out, H', W') |
权重: |
| 池化层 | nn.MaxPool2d(k) |
(B, c, H, W) |
(B, c, H/k, W/k) |
无参数 |
| 批量归一化 | nn.BatchNorm1d/2d |
同输入 | 同输入 | |
| Dropout | nn.Dropout(p) |
任意 | 相同 | 无参数（训练时随机遮罩） |
| LSTM | nn.LSTM(input, hidden) |
(seq, B, input) |
(seq, B, hidden) |
门控参数: |
| Transformer | nn.Transformer |
(seq, B, d_model) |
(seq, B, d_model) |
多头注意力+FFN |
以 nn.Linear(784, 256) 为例：
参数量 = 权重 + 偏置
= 784 × 256 + 256
= 200,704 + 256
= 200,960
以 nn.Conv2d(3, 64, kernel_size=3, padding=1) 为例：
参数量 = 权重: 3 × 64 × 3 × 3 + 偏置: 64
= 1,728 + 64
= 1,792
import torch.optim as optim
# 数据加载
train_loader = torch.utils.data.DataLoader(
dataset, batch_size=64, shuffle=True
)
# 模型、损失函数、优化器
model = SimpleCNN().to('cuda')
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练循环
for epoch in range(10):
running_loss = 0.0
for images, labels in train_loader:
images, labels = images.to('cuda'), labels.to('cuda')
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播
optimizer.zero_grad() # 清空梯度
loss.backward() # 计算梯度
optimizer.step() # 更新参数
running_loss += loss.item()
avg_loss = running_loss / len(train_loader)
print(f"Epoch [{epoch+1}/10], Loss: {avg_loss:.4f}")
| 任务类型 | 损失函数 | 说明 |
|---|---|---|
| 多分类 | nn.CrossEntropyLoss() |
内置 Softmax，输入 logits |
| 二分类 | nn.BCEWithLogitsLoss() |
内置 Sigmoid |
| 回归 | nn.MSELoss() |
均方误差 |
| 回归（鲁棒） | nn.L1Loss() |
平均绝对误差 |
| 对比学习 | nn.TripletMarginLoss() |
三元组损失 |
| 序列标注 | nn.CTCLoss() |
时序分类损失 |
| 优化器 | 特点 | 适合场景 | 收敛速度 | 内存占用 |
|---|---|---|---|---|
| SGD | 基本梯度下降 | 简单任务 | 慢 | 低 |
| SGD+Momentum | 加动量项加速 | 大多数 CV 任务 | 中 | 低 |
| Adam | 自适应学习率 | NLP、通用任务 | 快 | 高（保存动量） |
| AdamW | Adam + 解耦权重衰减 | 预训练、大模型 | 快 | 高 |
| RMSprop | 自适应学习率 | RNN、强化学习 | 中 | 中 |
| Adagrad | 自适应学习率 | 稀疏特征 | 中 | 累积梯度平方 |
学习率调度器：
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
# 或
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', patience=5, factor=0.1
)
Dataset 和 DataLoaderfrom torch.utils.data import Dataset, DataLoader
class CustomDataset(Dataset):
def __init__(self, data, labels, transform=None):
self.data = data
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
x = self.data[idx]
y = self.labels[idx]
if self.transform:
x = self.transform(x)
return x, y
# 使用 DataLoader 自动批处理和打乱
dataset = CustomDataset(data, labels, transform=transforms)
loader = DataLoader(dataset, batch_size=32, shuffle=True,
num_workers=4, pin_memory=True)
torchvision.transforms 数据增强| 变换 | 说明 | 训练效果 |
|---|---|---|
RandomHorizontalFlip() |
随机水平翻转 | 准确率提升 1-3% |
RandomRotation(degrees=15) |
随机旋转 ±15° | 泛化提升 1-2% |
ColorJitter(brightness=0.2) |
随机颜色抖动 | 鲁棒性提升 |
RandomResizedCrop(size=224) |
随机裁剪缩放 | 准确率提升 2-4% |
Normalize(mean, std) |
标准化 | 必须步骤 |
RandomErasing() |
随机擦除 | 相当于 Cutout |
AutoAugment() |
自动增强策略 | 最优效果，提升 2-5% |
典型组合：
from torchvision import transforms
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.2, 0.2, 0.2, 0.1),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")
if torch.cuda.is_available():
print(f"GPU型号: {torch.cuda.get_device_name(0)}")
print(f"显存容量: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
# 模型和数据移到 GPU
model = model.to(device)
data = data.to(device)
混合精度训练使用 torch.cuda.amp 在保持精度的同时减少显存占用约 50%，加速约 2-3 倍。
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler() # 梯度缩放，防止 underflow
for data, labels in loader:
data, labels = data.to(device), labels.to(device)
optimizer.zero_grad()
with autocast(): # 自动混合精度
outputs = model(data)
loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
精度与性能对比（ResNet-50, ImageNet, batch_size=256）：
| 精度模式 | 显存占用 | 训练速度 | Top-1 准确率 |
|---|---|---|---|
| FP32（全精度） | 10.2 GB | 1.0× | 76.15% |
| AMP（混合精度） | 5.8 GB | 2.3× | 76.13% |
| FP16（半精度，纯） | 5.1 GB | 2.5× | 75.82% |
最简单的方式，一个进程管理所有 GPU：
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])
缺点：主卡（GPU 0）负载过重，扩展性有限。
多进程方式，每个进程管理一个 GPU，效率更高：
# 启动脚本: torchrun --nproc_per_node=4 train.py
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])
# DataLoader 需要分布式采样器
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=32, sampler=sampler)
DDP vs DP 性能对比（4×A100, ResNet-50）：
| 方式 | 吞吐量（images/sec） | GPU 利用率 |
|---|---|---|
| 单卡 | 950 | 98% |
| DP 4卡 | 2,800 | 70% (主卡95%, 从卡55%) |
| DDP 4卡 | 3,720 | 97% |
| DDP 8卡 | 7,200 | 96% |
对大模型（>1B 参数）的分布式训练方案，将模型参数、梯度、优化器状态分片：
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(
model,
auto_wrap_policy=default_auto_wrap_policy,
mixed_precision=Bf16MixedPrecision()
)
FSDP 显存对比（LLaMA-7B 训练）：
| 配置 | 每卡显存 | 总显存 | 可训练的最大模型 |
|---|---|---|---|
| DDP (FP32) | 56 GB | 224 GB | ~3B |
| DDP (BF16) | 28 GB | 112 GB | ~7B |
| FSDP (BF16, 4卡) | 8.5 GB | 34 GB | ~13B |
将 PyTorch 模型转换为可序列化、可优化的表示：
# Tracing 方式（适合固定控制流）
traced_model = torch.jit.trace(model, example_input)
# Scripting 方式（适合动态控制流）
scripted_model = torch.jit.script(model)
# 保存和加载
traced_model.save('model.pt')
loaded = torch.jit.load('model.pt')
torch.compile（PyTorch 2.0）PyTorch 2.0 引入的编译优化，通过 TorchDynamo 捕获计算图，后端编译器（Inductor）优化执行：
# 一行代码加速
model = torch.compile(model, mode='reduce-overhead')
torch.compile 加速效果：
| 模式 | 说明 | ResNet-50 | BERT-Base | LLaMA-7B |
|---|---|---|---|---|
| 默认 | 平衡优化 | 1.4× | 1.3× | 1.2× |
reduce-overhead |
减少框架开销 | 1.8× | 1.5× | 1.3× |
max-autotune |
自动调优 | 2.0× | 1.7× | 1.5× |
跨框架模型交换格式：
torch.onnx.export(
model, dummy_input, "model.onnx",
input_names=['input'], output_names=['output'],
dynamic_axes={'input': {0: 'batch_size'},
'output': {0: 'batch_size'}}
)
| 量化方式 | API | 模型大小 | 推理速度 | 精度损失 |
|---|---|---|---|---|
| 动态量化 | torch.quantization.quantize_dynamic |
75% | 2-3× | <1% |
| 静态量化 | torch.quantization.quantize |
25% | 2-4× | 1-2% |
| QAT | torch.quantization.prepare_qat |
25% | 2-4× | <0.5% |
# 动态量化示例（推荐从这开始）
from torch.quantization import quantize_dynamic
quantized_model = quantize_dynamic(
model, {nn.Linear, nn.LSTM}, dtype=torch.qint8
)
# 模型大小降至 25%，推理速度提升 2-3 倍
torchvision 预训练模型import torchvision.models as models
# 所有预训练模型
available_models = {
'ResNet': models.resnet50(weights='DEFAULT'),
'EfficientNet': models.efficientnet_b0(weights='DEFAULT'),
'ViT': models.vit_b_16(weights='DEFAULT'),
'ConvNeXt': models.convnext_base(weights='DEFAULT'),
'Swin': models.swin_b(weights='DEFAULT'),
'DenseNet': models.densenet121(weights='DEFAULT'),
'MobileNet': models.mobilenet_v2(weights='DEFAULT'),
}
| 模型 | 参数量 | FLOPs | Top-1 | 推理延迟 (ms/batch) |
|---|---|---|---|---|
| ResNet-50 | 25.6M | 4.1G | 76.1% | 4.5 |
| ResNet-152 | 60.2M | 11.3G | 78.3% | 10.2 |
| EfficientNet-B0 | 5.3M | 0.4G | 77.1% | 2.8 |
| EfficientNet-B4 | 19.0M | 4.3G | 82.9% | 14.5 |
| EfficientNet-B7 | 66.0M | 39.0G | 84.3% | 52.0 |
| ViT-B/16 | 86.6M | 55.5G | 81.1% | 16.8 |
| ConvNeXt-Base | 88.6M | 15.4G | 84.0% | 18.2 |
| Swin-B | 87.8M | 15.4G | 83.5% | 20.1 |
| MobileNet-V2 | 3.5M | 0.3G | 71.9% | 1.5 |
测试环境：NVIDIA A100, batch_size=64, FP16
# 使用预训练模型进行迁移学习
model = models.resnet50(weights='DEFAULT')
# 替换最后一层（适配新任务，10类）
num_features = model.fc.in_features # 2048
model.fc = nn.Linear(num_features, 10)
# 可选：冻结特征提取层
for param in model.parameters():
param.requires_grad = False
for param in model.fc.parameters():
param.requires_grad = True # 只训练分类头
迁移学习效果（10% 数据量 vs 从头训练）：
| 场景 | 从头训练 | 冻结特征提取 | 全量微调 |
|---|---|---|---|
| 10% 数据 | 45.2% | 78.3% | 82.1% |
| 50% 数据 | 68.7% | 82.5% | 87.6% |
| 100% 数据 | 76.1% | 84.0% | 89.2% |
torchtext 和 NLP 支持class TransformerClassifier(nn.Module):
def __init__(self, vocab_size, d_model=256, nhead=8,
num_layers=4, num_classes=2):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.pos_encoder = nn.Parameter(
torch.randn(1, 512, d_model) * 0.1
)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=nhead, batch_first=True
)
self.transformer = nn.TransformerEncoder(
encoder_layer, num_layers=num_layers
)
self.classifier = nn.Linear(d_model, num_classes)
def forward(self, x):
# x: (batch, seq_len)
x = self.embedding(x) + self.pos_encoder[:, :x.size(1), :]
x = self.transformer(x)
x = x.mean(dim=1) # 池化
return self.classifier(x)
model = TransformerClassifier(vocab_size=30000)
print(f"参数量: {sum(p.numel() for p in model.parameters()):,}")
# 参数量: ~42M
以 d_model=768, nhead=12, num_layers=12（BERT-Base 级别）为例：
| 组件 | 计算公式 | 参数量 |
|---|---|---|
| Embedding | vocab_size × d_model |
30,000 × 768 = 23.0M |
| Self-Attention | 4 × d_model²（QKV+输出） |
4 × 768² = 2.36M |
| FFN | 2 × d_model × d_ff, d_ff=4d |
2 × 768 × 3072 = 4.72M |
| LayerNorm | 2 × d_model |
2 × 768 = 1.5K |
| 每层合计 | Self-Attention + FFN + LN | ~7.08M |
| 12层合计 | 12 × 7.08M | 84.96M |
| 总计 | Embedding + 12层 + 分类头 | ~110M |
from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter('runs/experiment_1')
# 记录标量
writer.add_scalar('Loss/train', loss, epoch)
writer.add_scalar('Accuracy/val', accuracy, epoch)
# 记录模型结构
writer.add_graph(model, dummy_input)
# 记录图片
writer.add_images('input_images', images[:16], epoch)
# 记录超参数
writer.add_hparams({'lr': 0.001, 'batch': 64}, {'hparam/acc': accuracy})
# 记录梯度直方图
for name, param in model.named_parameters():
writer.add_histogram(f'grad/{name}', param.grad, epoch)
writer.close()
| 工具 | 用途 |
|---|---|
torchinfo.summary(model, input_size) |
打印模型架构和参数量 |
torch.autograd.set_detect_anomaly(True) |
检测 NaN 梯度 |
torch.cuda.memory_summary() |
分析 GPU 显存使用 |
torch.cuda.empty_cache() |
清理缓存显存 |
torch.jit.trace(model, input) |
检查模型输出形状 |
torch.utils.bottleneck |
性能瓶颈分析 |
torch.profiler |
详细的性能分析 |
| 错误 | 原因 | 解决方案 |
|---|---|---|
RuntimeError: CUDA out of memory |
显存不足 | 减小 batch_size，使用 AMP、梯度累积 |
gradient computation has been modified |
inplace 操作破坏计算图 | 禁用 inplace=True 或使用 x = x + 1 代替 x += 1 |
device-side assert triggered |
标签超出范围 | 检查 labels 是否在 [0, num_classes) 内 |
Expected all tensors on same device |
混合 CPU/GPU 计算 | 使用 .to(device) 统一设备 |
NaN in loss |
梯度爆炸 | 梯度裁剪 (clip_grad_norm_)、降低学习率 |
multi-target not supported |
CrossEntropyLabel 形状错误 | 确保 labels 是 1D，非 one-hot |
| 库名 | 用途 |
|---|---|
| torchvision | 计算机视觉：预训练模型、数据集、图像变换 |
| torchtext | 自然语言处理：数据集、分词、词向量 |
| torchaudio | 音频处理：数据集、特征提取、变换 |
| torchrec | 推荐系统：大规模嵌入、分布式训练 |
| torchdata | 数据流水线：高效数据加载和预处理 |
| torchx | 分布式训练编排：SLURM、Kubernetes |
| torchserve | 模型部署：REST API 服务 |
| torchao | 量化/剪枝/蒸馏：模型压缩 |
| torchft | 容错训练：检查点和故障恢复 |
| 类别 | 推荐库 |
|---|---|
| 实验管理 | MLflow, Weights & Biases, Neptune |
| 超参搜索 | Optuna, Ray Tune |
| 视觉工具 | OpenMMLab (mmdetection, mmsegmentation) |
| NLP 框架 | Hugging Face Transformers, spaCy |
| 强化学习 | Stable-Baselines3, RLlib |
| 图神经网络 | PyTorch Geometric |
| 概率编程 | Pyro, Bean Machine |
| 模型压缩 | Intel Neural Compressor, NNI |
| 推理优化 | TensorRT, ONNX Runtime, vLLM |
project/
├── config.py # 超参数配置
├── data/
│ ├── dataset.py # 数据集类
│ └── transforms.py # 数据增强
├── models/
│ └── architecture.py # 模型定义
├── trainer.py # 训练逻辑
├── evaluate.py # 评估逻辑
├── utils/
│ ├── logger.py # 日志
│ └── metrics.py # 评价指标
├── main.py # 入口
└── requirements.txt # 依赖
max_norm=1.0）nn.CrossEntropyLoss(label_smoothing=0.1)）| 优化手段 | 加速效果 | 实现复杂度 |
|---|---|---|
DataLoader(num_workers=4, pin_memory=True) |
1.2× | 低 |
| 混合精度 (AMP) | 2.0× | 低 |
torch.compile |
1.5× | 极低 |
| 梯度累积 | 显存受限时有效 | 低 |
| DDP 分布式训练 | 4卡 3×, 8卡 6× | 中 |
使用 channels_last 内存格式 |
1.1× | 低 |
set_benchmark(True) |
1.05× | 极低 |
| 编译 CUDA 自定义算子 | 2-10× | 高 |
使用 torch.jit.script |
1.2× | 中 |
参见：