如何在PyTorch中为特定张量释放GPU内存?

How to free GPU memory for a specific tensor in PyTorch?

我目前正在使用PyTorch运行深度学习程序,并希望为特定张量释放GPU内存。

我想到了deltorch.cuda.empty_cache()之类的方法,但是del似乎无法正常工作(我什至不确定是否完全释放了内存),并且torch.cuda.empty_cache()似乎可以释放所有未使用的内存 ,但是我只想释放特定张量的内存。

PyTorch中有提供此功能的功能吗?

提前致谢。


del运算符可以工作,但是您不会看到所使用的GPU内存减少,因为该内存未返回到cuda设备。这是一种优化技术,从用户的角度来看,内存已被"释放"。即,该存储器现在可用于制作新的张量。

资料来源:Pytorch论坛


如果以后要执行自动微分,PyTorch会将张量存储在计算图中(如果使用requires_grad = True进行了初始化)。如果您不想再使用特定的张量进行梯度计算,则可以使用detach方法告诉PyTorch它不再需要存储该张量的值用于梯度计算。这将有助于释放一些内存(仅删除该特定张量而不删除整个计算图)。

例如-my_tensor.detach()


obj = Nonedel obj都相似,除了del将删除引用。

但是,您需要调用gc.collect()来释放Python内存,而无需重新启动笔记本计算机。

如果您想从PyTorch缓存中清除obj,请运行:

1
torch.cuda.empty_cache()

在执行最后一条命令后,Nvidea smi或nvtop将注意到您的操作。