pytorch

PyTorch - 深度学习与张量

PyTorch 是一个基于 Python 的科学计算包，利用图形处理器（GPU）的强大性能，并通过动态计算图系统提供最大的灵活性和速度。

适用场景

• 构建和训练深度神经网络（CNN、RNN、Transformers）。
• 研究需要动态计算图的新型 AI 架构。
• 在 NVIDIA（CUDA）或 Mac（MPS）硬件上加速张量运算。
• 求解物理信息神经网络（PINN）。
• 实现生成模型（GAN、Diffusion）。
• 使用 Autograd（自动微分）进行大规模优化。
• 部署生产级 AI（通过 TorchScript/ONNX）。

参考文档

官方文档：https://pytorch.org/docs/
教程：https://pytorch.org/tutorials/
搜索模式：torch.nn、torch.optim、torch.utils.data、Autograd、Tensor.to(device)

核心原则

张量

核心数据结构，类似于 NumPy 的 ndarray，但有两个关键的扩展：它可以存储在 GPU 上，并支持自动微分。

动态计算图（Autograd）

PyTorch 会在代码执行时即时构建计算图。这样，模型中就可以使用标准的 Python 控制流（if/for）。

模块和参数

nn.Module 是所有神经网络组件的基类。它会自动追踪用于优化的 nn.Parameter 对象（权重/偏置）。

快速参考

安装

# CPU
pip install torch torchvision
# GPU (Check pytorch.org for specific CUDA versions)
pip install torch --index-url https://download.pytorch.org/whl/cu121

标准导入

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

基本模式 - 简单线性回归（“PyTorch 方式”）

import torch

# 1. Data (Tensors)
X = torch.tensor([[1.0], [2.0], [3.0]], requires_grad=True)
y = torch.tensor([[2.0], [4.0], [6.0]])

# 2. Simple Model
model = torch.nn.Linear(1, 1) # y = w*x + b

# 3. Loss and Optimizer
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 4. Training Loop
for epoch in range(100):
    prediction = model(X)
    loss = criterion(prediction, y)
    
    optimizer.zero_grad() # Clear previous gradients
    loss.backward()       # Compute gradients (Autograd)
    optimizer.step()      # Update weights

关键规则

✅ 应该做

• 始终调用 zero_grad() - PyTorch 默认会累积梯度。忘记调用会导致模型无法收敛。
• 使用 .to(device) - 明确将模型和数据移动到同一设备（CPU 或 CUDA）。
• 使用 DataLoader - 不要在循环中手动输入数据；DataLoader 会处理批处理、打乱数据和多进程加载。
• 设置 model.train() / model.eval() - 对于 Dropout 和 BatchNorm 等在推理期间行为不同的层，这一点至关重要。
• 推理时使用 torch.no_grad() - 不构建计算图可以显著节省内存和计算量。
• 指定 dtypes - 注意 float32（标准精度）、float64（科学计算精度）和 float16（速度/GPU）。

❌ 不应该做

• 混用 CPU 和 GPU 张量 - RuntimeError: Expected all tensors to be on the same device 是最常见的错误。
• 使用标准 Python 循环进行数学运算 - 使用向量化张量运算以提升性能。
• 忘记 .item() - 从张量获取标量值用于日志记录时，使用 loss.item() 将其从计算图中分离。
• 在 GPU 上过度使用 float64 - 许多消费级 GPU 的双精度性能较差；尽可能使用 float32。

反模式（切勿使用）

import torch

# ❌ BAD: Mixing Python lists/arrays with Tensors in a loop
# for x in data:
#     res = model(torch.tensor(x)) # Extremely slow re-allocation!

# ✅ GOOD: Batching
# data_tensor = torch.stack([torch.tensor(x) for x in data])
# res = model(data_tensor)

# ❌ BAD: Calculating loss without zeroing gradients
loss.backward()
optimizer.step()
# Next iteration... gradients will be double what they should be!

# ✅ GOOD:
optimizer.zero_grad()
loss.backward()
optimizer.step()

# ❌ BAD: Standard NumPy for inference
# with torch.no_grad():
#    pred = model(X).numpy() # Can be slow on GPU if not handled

# ✅ GOOD: Explicit move to CPU
# pred = model(X).detach().cpu().numpy()

张量和设备管理

在 CPU 和 GPU 之间移动

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Create tensor on device
x = torch.randn(3, 3, device=device)

# Move model to device
model = MyModel().to(device)

# Move data to device during loop
for inputs, labels in dataloader:
    inputs, labels = inputs.to(device), labels.to(device)
    # ...

构建模型（nn.Module）

灵活的架构

class ScientificNet(nn.Module):
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, 1)
        self.dropout = nn.Dropout(0.2)
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = self.dropout(x)
        x = torch.sigmoid(self.layer2(x))
        return x

model = ScientificNet(10, 50)

自定义数据集（torch.utils.data）

处理科学数据文件（例如 HDF5 或 CSV）

class MyScientificDataset(Dataset):
    def __init__(self, file_path):
        self.data = pd.read_csv(file_path)
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        # Convert row to tensor
        sample = torch.tensor(self.data.iloc[idx, :-1].values, dtype=torch.float32)
        label = torch.tensor(self.data.iloc[idx, -1], dtype=torch.float32)
        return sample, label

dataset = MyScientificDataset("experiment_results.csv")
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

高级 Autograd

用于物理计算的梯度（Jacobian/Hessian）

x = torch.linspace(-5, 5, 100, requires_grad=True)
y = x**3

# First derivative dy/dx
# create_graph=True allows for higher-order derivatives
dy_dx = torch.autograd.grad(y.sum(), x, create_graph=True)[0]

# Second derivative (Hessian) d2y/dx2
d2y_dx2 = torch.autograd.grad(dy_dx.sum(), x)[0]

实用工作流

1. 物理信息神经网络（PINN）片段

def pde_loss(model, x):
    """Simple ODE: u'(x) = u(x)."""
    x.requires_grad = True
    u = model(x)
    u_x = torch.autograd.grad(u.sum(), x, create_graph=True)[0]
    return F.mse_loss(u_x, u)

# Training loop combines data_loss + pde_loss

2. 科学训练中的提前停止

best_loss = float('inf')
patience = 10
counter = 0

for epoch in range(1000):
    train_loss = train_one_epoch()
    val_loss = validate()
    
    if val_loss < best_loss:
        best_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pth')
        counter = 0
    else:
        counter += 1
        if counter >= patience:
            print("Early stopping triggered")
            break

3. 用于化学的特征提取

def extract_embeddings(model, loader):
    model.eval()
    embeddings = []
    with torch.no_grad():
        for batch in loader:
            # Assume model has a .get_features() method
            features = model.get_features(batch.to(device))
            embeddings.append(features.cpu())
    return torch.cat(embeddings)

性能优化

使用 torch.compile（PyTorch 2.0+）

只需一行代码，即可显著加速现代模型：

model = MyModel()
compiled_model = torch.compile(model)

混合精度（torch.cuda.amp）

节省内存并加快现代 GPU（Tensor Core）上的训练速度。

scaler = torch.cuda.amp.GradScaler()

for inputs, labels in loader:
    with torch.cuda.amp.autocast():
        output = model(inputs)
        loss = criterion(output, labels)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

常见问题及解决方案

梯度消失/爆炸

# ❌ Problem: Loss becomes 'nan' or weights don't update
# ✅ Solution: Use Gradient Clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

ReLU 神经元死亡

如果神经元的输出始终 <= 0，它就会停止学习。

# ✅ Solution: Use LeakyReLU or ELU
self.act = nn.LeakyReLU(0.01)

内存泄漏（张量滞留在计算图中）

直接记录 loss 会使整个计算图一直保留在内存中。

# ❌ BAD: total_loss += loss
# ✅ GOOD: total_loss += loss.item()

PyTorch 是 AI 革命的引擎。对于科学家来说，它提供了从传统数据分析通往可微模型世界的桥梁，让人们能够发现此前无法察觉的模式。