PyTorch 开发模式

构建稳健、高效和可复现深度学习应用的 PyTorch 惯用模式与最佳实践。

何时使用

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

核心原则

1. 设备无关代码

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

# Good: Device-agnostic
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

# Bad: Hardcoded device
model = MyModel().cuda()  # Crashes if no GPU
data = data.cuda()

2. 可复现性优先

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

# Good: Full reproducibility setup
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

# Bad: No seed control
model = MyModel()  # Different weights every run

3. 显式形状管理

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

# Good: Shape-annotated forward pass
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…