Pytorch模型训练CPU内存泄漏问题

Question

当我在GPU设备上训练我的 pytorch 模型时，我的 python 脚本突然被杀死。深入研究操作系统日志文件，我发现脚本被 OOM Killer 杀死，因为我的CPU内存不足。It\xe2\x80\x99s 非常奇怪的是，我在GPU设备上训练了模型，但CPU内存用完了。\n OOM 杀手日志文件的快照\n

\n

为了调试这个问题，我安装了 python 内存分析器。从内存分析器查看日志文件，我发现当发生列-=操作时，我的 CPU 内存逐渐增加，直到 OOM Killer 杀死我的\n程序。\n Python 内存分析器的快照\n\n这\xe2\x80\x99s很奇怪，我尝试了很多方法来解决这个问题。最后，我发现在赋值操作之前，我先分离Tensor。神奇的是，它解决了这个问题。但我不明白\xe2\x80\x99不明白清楚为什么它有效。这是我原来的函数代码。

\n

def GeneralizedNabla(self, image):\n        pad_size = 2\n        affinity = torch.zeros(image.shape[0], self.window_size**2, self.h, self.w).to(self.device)\n        h = self.h+pad_size\n        w = self.w+pad_size\n        #pad = nn.ZeroPad2d(pad_size)\n        image_pad = self.pad(image)\n        for i in range(0, self.window_size**2):\n            affinity[:, i, :, :] = image[:, :, :].detach()  # initialization\n            dy = int(i/5)-2\n            dx = int(i % 5)-2\n            h_start = pad_size+dy\n            h_end = h+dy  # if 0 <= dy else h+dy\n            w_start = pad_size+dx\n            w_end = w+dx  # if 0 <= dx else w+dx\n            affinity[:, i, :, :] -= image_pad[:, h_start:h_end, w_start:w_end].detach()\n        self.Nabla=affinity\n        return\n

\n

如果大家有什么想法，我将不胜感激，谢谢。

\n

Answer 1

以前，当你没有.detach()在张量上使用 时，你也在累积计算图，并且随着你的继续，你不断累积越来越多，直到你最终耗尽你的记忆直至崩溃。
当您执行 a 时detach()，您可以有效地获取数据，而无需计算梯度所需的先前纠缠的历史记录。