基于MFDNet的Python视频逐帧去雨显存溢出问题排查

核心原因 #

• 原代码里的restored = restored.cpu().detach().numpy()是关键的显存释放步骤：它把GPU上的张量移到CPU、从计算图彻底分离、转为numpy数组，这会直接释放该张量占用的GPU显存，同时切断梯度依赖。
• 移除这段代码后，推理生成的restored张量会一直留在GPU显存中，而且PyTorch默认会保留张量的计算图关联（哪怕是推理模式）。逐帧处理视频时，每帧的restored都会在GPU里累积，最终撑爆显存。

解决办法 #

方案1：保留显存释放逻辑（无需转numpy） #

如果不需要转成numpy数组，只需要释放GPU显存，替换成以下代码：

# 仅完成显存释放操作，不转numpy
restored = restored.detach().cpu()
# 手动触发垃圾回收，加速显存释放
import gc
gc.collect()
torch.cuda.empty_cache()

方案2：关闭梯度计算减少显存占用 #

在推理代码外层加torch.no_grad()上下文，阻止计算图生成，从根源减少显存消耗：

with torch.no_grad():
    restored = model(input_frame)
    restored = torch.clamp(restored, 0, 1)
    # 这里写后续帧处理逻辑

方案3：逐帧手动清理张量 #

每帧处理完成后，显式删除不再需要的张量，强制释放显存：

# 当前帧处理完毕后执行
del restored
import gc
gc.collect()
torch.cuda.empty_cache()

额外提醒 #

• 用torch.save保存模型或结果时，要先把张量移到CPU再保存，避免GPU张量的显存占用无法释放。
• 视频逐帧处理是长循环，PyTorch自动垃圾回收可能不及时，手动调用gc.collect()和torch.cuda.empty_cache()能确保显存及时释放。

内容的提问来源于stack exchange，提问作者JimmyHu