Question

我在this file中找到了一行代码：

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

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

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

Answer 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反向传播的“无”渐变。

Answer 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

如何仅在正向传递中将pytorch模块用作纯特征提取器？

2 个答案: