每次转发后如何释放临时消耗的GPU内存？

Question

我有这样的课程：

class Stem(nn.Module):

    def __init__(self):
        super(Stem, self).__init__()
        self.out_1 = BasicConv2D(3, 32, kernelSize = 3, stride = 2)
        self.out_2 = BasicConv2D(32, 32, kernelSize = 3, stride = 1)
        self.out_3 = BasicConv2D(32, 64, kernelSize = 3, stride = 1, padding = 1)


    def forward(self, x):
        x = self.out_1(x)
        x = self.out_2(x)
        x = self.out_3(x)

        return x

而Stem的属性out_1,2,3是以下类的实例：

^{pr2}$

当训练时，在Stem.forward()内，nvidia-smi告诉每一行将消耗xMBs GPU内存，但在Stem.forward()完成后，内存不会释放，导致训练很快崩溃，耗尽GPU内存。在

因此，问题是： 如何释放临时占用的GPU内存？在

Answer 1

您的模型看起来确实不错，所以您可能需要大致了解一下pytorch是如何管理内存分配的。我怀疑你只是简单地保持指向你的返回值（y）的指针（例如，通过累积损失或类似的方法）。由于Pythorch存储了整个附加的计算图，因此永远不会释放内存。在

有关更详细的讨论，请参见this question，尤其是{a2}。在