PyTorch详解

PyTorch详解

PyTorch 设计哲学

动态计算图（Dynamic Computation Graph）

计算图在运行时动态构建（称为 “Define-by-Run”），允许灵活修改模型结构，适合研究和新算法实验。

对比：TensorFlow 1.x 的静态图需先定义后执行，而 PyTorch 的动态图更接近 Python 原生编程体验。

Python 原生风格

深度集成 Python 生态，支持直接使用 if、for 等控制流，调试方便（如 pdb）。

GPU 优先的张量计算

torch.Tensor 提供类似 NumPy 的接口，但支持自动微分和 GPU 加速。

核心组件与功能

张量（Tensor）

基础操作：

import torch
x = torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)  # 创建张量
y = torch.randn(2, 2, device='cuda')                   # GPU 张量
z = x + y.to('cpu')                                    # 跨设备运算

自动微分（Autograd）：

x = torch.tensor(3.0, requires_grad=True)
y = x**2 + 2*x
y.backward()  # 计算梯度
print(x.grad) # 输出: tensor(8.0)  (dy/dx = 2x + 2)

神经网络模块（torch.nn）

定义模型：

class CNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=3)
        self.fc = torch.nn.Linear(32 * 26 * 26, 10)  # MNIST 分类

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = x.view(-1, 32 * 26 * 26)  # 展平
        return self.fc(x)
model = CNN().cuda()  # 移至 GPU

损失函数与优化器：

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

数据处理（torch.utils.data）

自定义数据集：

from torch.utils.data import Dataset, DataLoader
class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.data[idx]

dataset = MyDataset(torch.randn(100, 3))
loader = DataLoader(dataset, batch_size=16, shuffle=True)

训练循环

for epoch in range(10):
    for inputs, labels in loader:
        optimizer.zero_grad()
        outputs = model(inputs.cuda())
        loss = criterion(outputs, labels.cuda())
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

高级特性

动态图控制流

def dynamic_forward(x, n_layers):
    for i in range(n_layers):
        if i % 2 == 0:
            x = torch.relu(x)
        else:
            x = torch.sigmoid(x)
    return x

混合精度训练（AMP）

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

分布式训练

单机多卡（DataParallel / DistributedDataParallel）：

model = torch.nn.DataParallel(model)  # 简单并行
# 或使用更高效的 DDP
torch.distributed.init_process_group(backend='nccl')
model = torch.nn.parallel.DistributedDataParallel(model)

模型部署与生产化

TorchScript（模型序列化）

scripted_model = torch.jit.script(model)  # 转换为脚本
scripted_model.save("model.pt")
# 加载模型
loaded_model = torch.jit.load("model.pt")

ONNX 导出

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=11)

移动端部署（LibTorch）

使用 C++ 加载 PyTorch 模型：

torch::jit::script::Module module = torch::jit::load("model.pt");
torch::Tensor input = torch::ones({1, 3, 224, 224});
torch::Tensor output = module.forward({input}).toTensor();

生态工具

工具库	用途
TorchVision	计算机视觉（数据集、预训练模型）
TorchText NLP	文本处理
PyTorch Lightning	简化训练流程的高层封装
HuggingFace Transformers	预训练 NLP 模型（BERT、GPT等）
Fast.ai	高级 API 快速训练模型

与 TensorFlow 的关键对比

特性	PyTorch	TensorFlow
计算图	动态图（即时执行）	静态图（TF 1.x）或动态图（TF 2.x）
调试	直接使用 Python 调试工具	依赖 TensorBoard 或 Eager 模式
部署	需转换为 TorchScript/ONNX	原生支持（TF Serving、TFLite）
社区	学术界主导，研究前沿	工业界广泛使用
API	设计 灵活，接近原生 Python	Keras 标准化，适合快速开发

经典应用示例

Transformer 模型（NLP）

from transformers import BertModel
model = BertModel.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hello, PyTorch!", return_tensors="pt")
outputs = model(**inputs)

GAN（生成对抗网络）

generator = nn.Sequential(
    nn.Linear(100, 256),
    nn.LeakyReLU(),
    nn.Linear(256, 784),
    nn.Tanh()
)
discriminator = nn.Sequential(
    nn.Linear(784, 256),
    nn.LeakyReLU(),
    nn.Linear(256, 1),
    nn.Sigmoid()
)

适用场景建议

选择 PyTorch 若：

需要快速实验或研究新模型。

重视代码可读性和调试便利性。

从事学术界或前沿 AI 研究（如 NLP、GAN）。

选择 TensorFlow 若：

需要直接部署到生产环境（如移动端、Web）。

使用 Keras 快速构建标准模型（如图像分类）。