在pytorch中训练后如何清除显卡内存？

Question

我在 colab 中处理 pytorch
pytorch 训练时消耗大量内存
训练结束后，我保存模型，并将模型加载到另一个笔记本（注2）。
在注释 2 中，在加载 state_dict 和所有内容之后，pytorch 消耗的内存比在训练状态下少得多。
所以，我想知道 'useless' 数据在训练时存储在显卡内存中（在我的例子中，大约 13gb）...
如果是这样，如何在训练后删除无用数据？

加上。我试图删除训练时使用的变量，但不够大（大约 2gb）

Answer 1

这在训练中是意料之中的。在训练过程中，操作本身会占用内存。

例如，考虑以下操作 -

a = np.random.rand(100, 500, 300)
b = np.random.rand(200, 500, 300)
c = (a[:, None, :, :] * b[None, :, :, :]).sum(-1).sum(-1)

a、b、c的内存大小分别在400MB左右。但是，如果您检查

%memit (a[:, None, :, :] * b[None, :, :, :]).sum(-1).sum(-1)

那是 23 GB！该行本身占用大量内存来实际执行操作，因为涉及大量中间数组。这些数组是临时的，操作结束后会自动删除。所以你删除一些变量不会对减少足迹有多大作用。

解决这个问题的方法是使用内存优化操作。

例如，执行 np.tensordot(a, b, ((1, 2), (1, 2))) 而不是通过广播进行乘法会留下更好的内存占用空间。

所以你需要做的是确定代码中的哪个操作需要如此大的内存，看看你是否可以用内存效率更高的等效项替换它（根据你的具体用途，这甚至可能是不可能的-case).

在pytorch中训练后如何清除显卡内存？

how can I clear graphic card memory after training in pytorch?

memory-leaks

deep-learning

pytorch