如何在不计算梯度的情况下使用 PyTorch 模型？答案

【问题标题】：How to use a PyTorch model without calculating gradients?如何在不计算梯度的情况下使用 PyTorch 模型？
【发布时间】：2019-04-02 08:42:46
【问题描述】：

在this file找到一行代码：

out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

如您所见forward()函数在backward()之前被多次调用，在我的测试中，GPU消耗会相应增加，因此会发生GPU内存泄漏。这里的问题是：

如何解决这个问题？
是否可以将模型用作这三个连续调用的纯特征提取器（当然要进行一些修改），然后在第四次调用时将其用作forward()，如果“是”，如何实现？

【问题讨论】：

我是否正确地理解了您只想计算 out_X 值，然后在第四次调用中使用它们进行反向计算？您是否已经研究过model.eval()（我认为较新的 PyTorch 版本可能会有不同的模式）？

标签： memory-leaks gpu pytorch

【解决方案1】：

如果您只想将模型作为纯特征提取器（即过程中不更新），您可以这样做

假设您有一个模型，model，您想先将其设置为评估模式。

model = FeatureExtractor().cuda()
mode.eval()

接下来，为了减少gpu内存的消耗，你需要停止梯度累积和计算。

for p in model.parameters():
    p.require_grads = False
out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

那么，我想你的问题就解决了。请注意，您可以通过out_a、out_p 和out_n 反向传播“无”渐变。

【讨论】：

【解决方案2】：

作为已经提出的答案的替代方案，您还可以使用 torch.no_grad() 在本地禁用梯度计算：

with torch.no_grad():
    out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

这还可以防止with 语句中的任何梯度计算，从而显着减少内存使用：

禁用梯度计算的上下文管理器。

禁用梯度计算对于推理很有用，当您确保您不会致电Tensor.backward()。会减少记忆否则会有的计算消耗 requires_grad=True。在这种模式下，每次计算的结果都会有 requires_grad=False，即使输入有 requires_grad=True。

来源：https://pytorch.org/docs/stable/autograd.html#torch.autograd.no_grad

【讨论】：