PyTorch 模式指南

PyTorch 开发模式

用于构建健壮、高效且可复现的深度学习应用程序的惯用 PyTorch 模式和最佳实践。

何时激活

• 编写新的 PyTorch 模型或训练脚本
• 审查深度学习代码
• 调试训练循环或数据管道
• 优化 GPU 内存使用或训练速度
• 设置可复现的实验

核心原则

1. 与设备无关的代码

始终编写能在 CPU 和 GPU 上运行且不硬编码设备的代码。

# 推荐：与设备无关
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

# 不推荐：硬编码设备
model = MyModel().cuda()  # 如果没有 GPU 会崩溃
data = data.cuda()

2. 可复现性优先

设置所有随机种子以获得可复现的结果。

# 推荐：完整的可复现性设置
def set_seed(seed: int = 42) -> None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 不推荐：无种子控制
model = MyModel()  # 每次运行权重都不同

3. 显式形状管理

始终记录和验证张量形状。

# 推荐：带形状注释的前向传播
def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (batch_size, channels, height, width)
    x = self.conv1(x)    # -> (batch_size, 32, H, W)
    x = self.pool(x)     # -> (batch_size, 32, H//2, W//2)
    x = x.view(x.size(0), -1)  # -> (batch_size, 32*H//2*W//2)
    return self.fc(x)    # -> (batch_size, num_classes)

# 不推荐：无形状跟踪
def forward(self, x):
    x = self.conv1(x)
    x = self.pool(x)
    x = x.view(x.size(0), -1)  # 这到底是什么尺寸？
    return self.fc(x)           # 这能正常工作吗？

模型架构模式

清晰的 nn.Module 结构

# 推荐：组织良好的模块
class ImageClassifier(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64 * 16 * 16, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# 不推荐：所有事情都在 forward 中
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = F.conv2d(x, weight=self.make_weight())  # 每次调用都创建权重！
        return x

正确的权重初始化

# 推荐：显式初始化
def _init_weights(self, module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.BatchNorm2d):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model = MyModel()
model.apply(model._init_weights)

训练循环模式

标准训练循环

# 推荐：包含最佳实践的完整训练循环
def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    scaler: torch.amp.GradScaler | None = None,
) -> float:
    model.train()  # 始终设置训练模式
    total_loss = 0.0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad(set_to_none=True)  # 比 zero_grad() 更高效

        # 混合精度训练
        with torch.amp.autocast("cuda", enabled=scaler is not None):
            output = model(data)
            loss = criterion(output, target)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

验证循环

# 推荐：正确评估
@torch.no_grad()  # 比包装在 torch.no_grad() 块中更高效
def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()  # 始终设置评估模式——禁用 dropout，使用运行中的 BN 统计信息
    total_loss = 0.0
    correct = 0
    total = 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        total_loss += criterion(output, target).item()
        correct += (output.argmax(1) == target).sum().item()
        total += target.size(0)

    return total_loss / len(dataloader), correct / total

数据管道模式

自定义数据集

# 推荐：带类型提示的清晰数据集
class ImageDataset(Dataset):
    def __init__(
        self,
        image_dir: str,
        labels: dict[str, int],
        transform: transforms.Compose | None = None,
    ) -> None:
        self.image_paths = list(Path(image_dir).glob("*.jpg"))
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        img = Image.open(self.image_paths[idx]).convert("RGB")
        label = self.labels[self.image_paths[idx].stem]

        if self.transform:
            img = self.transform(img)

        return img, label

高效的数据加载器配置

# 推荐：优化后的 DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # 训练时打乱
    num_workers=4,           # 并行数据加载
    pin_memory=True,         # 更快的 CPU->GPU 传输
    persistent_workers=True, # 在 epoch 之间保持工作进程存活
    drop_last=True,          # 为 BatchNorm 保持一致的批次大小
)

# 不推荐：缓慢的默认设置
dataloader = DataLoader(dataset, batch_size=32)  # num_workers=0, no pin_memory

变长数据的自定义 Collate

# 推荐：在 collate_fn 中填充序列
def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
    sequences, labels = zip(*batch)
    # 填充到批次中的最大长度
    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels)

dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

检查点模式

保存和加载检查点

# 推荐：包含所有训练状态的完整检查点
def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: str,
) -> None:
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, path)

def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint

# 不推荐：只保存模型权重（无法恢复训练）
torch.save(model.state_dict(), "model.pt")

性能优化

混合精度训练

# 推荐：带 GradScaler 的 AMP
scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
    with torch.amp.autocast("cuda"):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

大型模型的梯度检查点

# 推荐：用计算换内存
from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 在反向传播期间重新计算激活以节省内存
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

用于加速的 torch.compile

# 推荐：编译模型以加快执行速度（PyTorch 2.0+）
model = MyModel().to(device)
model = torch.compile(model, mode="reduce-overhead")

# 模式："default" (安全), "reduce-overhead" (更快), "max-autotune" (最快)

快速参考：PyTorch 惯用法

惯用法	描述
model.train() / model.eval()	在训练/评估前始终设置模式
torch.no_grad()	推理时禁用梯度
optimizer.zero_grad(set_to_none=True)	更高效的梯度清零
.to(device)	与设备无关的张量/模型放置
torch.amp.autocast	混合精度，速度提升 2 倍
pin_memory=True	更快的 CPU→GPU 数据传输
torch.compile	用于加速的 JIT 编译 (2.0+)
weights_only=True	安全模型加载
torch.manual_seed	可复现的实验
gradient_checkpointing	用计算换内存

应避免的反模式

# 不推荐：验证时忘记 model.eval()
model.train()
with torch.no_grad():
    output = model(val_data)  # Dropout 仍处于活动状态！BatchNorm 使用批次统计信息！

# 推荐：始终设置评估模式
model.eval()
with torch.no_grad():
    output = model(val_data)

# 不推荐：破坏自动求导的原地操作
x = F.relu(x, inplace=True)  # 可能破坏梯度计算
x += residual                  # 原地加法破坏自动求导图

# 推荐：非原地操作
x = F.relu(x)
x = x + residual

# 不推荐：在训练循环中重复将数据移动到 GPU
for data, target in dataloader:
    model = model.cuda()  # 每次迭代都移动模型！

# 推荐：在循环前将模型移动一次
model = model.to(device)
for data, target in dataloader:
    data, target = data.to(device), target.to(device)

# 不推荐：在 backward 之前使用 .item()
loss = criterion(output, target).item()  # 从图中分离！
loss.backward()  # 错误：无法通过 .item() 反向传播

# 推荐：仅用于日志记录时调用 .item()
loss = criterion(output, target)
loss.backward()
print(f"Loss: {loss.item():.4f}")  # 在 backward 之后使用 .item() 没问题

# 不推荐：不正确使用 torch.save
torch.save(model, "model.pt")  # 保存整个模型（脆弱，不可移植）

# 推荐：保存 state_dict
torch.save(model.state_dict(), "model.pt")

记住：PyTorch 代码应与设备无关、可复现且注意内存使用。如有疑问，请使用 torch.profiler 进行分析，并使用 torch.cuda.memory_summary() 检查 GPU 内存。