Pytorch-Tricks

Tue, 07 Apr 2026 00:00:00 +0000

写在前面：Pytorch 内容繁杂，实验中用到的技巧与机制若不及时记录，容易遗忘并导致重复造轮子。本文档用于记录 Pytorch 的相关技巧与机制，将持续更新，方便后续查阅。

动态计算图

在一次实验日志中，发现一个固定的 CNN 模块推理时间存在较大波动。排查后发现是数据流水线卡顿或 GPU 内存管理抖动所致，但也因此注意到了 Pytorch 的动态计算图机制。

Pytorch 在每次前向传播时都会重新构建计算图，这意味着每次前向传播都会产生额外开销，尤其在模型结构复杂或输入数据变化较大时，可能导致推理时间波动。

具体而言：

模型中包含数据依赖的条件分支（如 if 语句）时，不同前向传播可能执行不同计算路径，导致计算图结构变化。

class BadEncoder(nn.Module):
 def forward(self, x):
 if x.sum() > 0: # 数据依赖的条件
 x = self.heavy_conv(x)
 else:
 x = self.light_conv(x)
 return x

处理变长序列（如 RNN 或 Transformer）时，每次输入长度可能不同，导致计算图结构变化。

class BadRNN(nn.Module):
 def forward(self, x):
 for i in range(x.size(1)): # 依赖输入序列长度
 x = self.rnn_cell(x[:, i, :], x)
 return x

修改模型与训练代码后发现速度明显下降。排查发现数据加载耗时过长，遂将 DataLoader 的在线预处理改为离线处理。但模型推理速度仍慢于基线，后续需进一步优化。