PyTorch：释放CUDA内存的秘诀

PyTorch 释放 CUDA
PyTorch是一个广泛使用的深度学习框架，它支持使用GPU进行加速计算。在使用PyTorch进行深度学习训练或推理时，我们通常会使用CUDA（Compute Unified Device Architecture）来将计算任务从CPU转移到GPU上。然而，一旦我们完成了计算任务，最好将CUDA内存释放回系统，以便其他应用程序可以使用它。
在PyTorch中，释放CUDA内存非常简单。首先，您需要确保已经完成了所有使用CUDA内存的任务。然后，您可以采取以下措施释放CUDA内存：

1. 将变量设置为 None。这将触发垃圾收集器将其从内存中删除。

   import torch
   # 创建一个张量并将其分配给GPU
   x = torch.randn(100).cuda()
   # 现在将张量设置为None以释放内存
   x = None

2. 使用 torch.cuda.empty_cache() 函数。这将释放当前未使用的所有CUDA内存。这对于确保您的应用程序不会耗尽所有可用的CUDA内存非常有用。但是请注意，此操作并不能保证内存的绝对释放，因为有可能是其他部分的代码在使用这些内存。

   import torch
   # 创建一些张量并将其分配给GPU
   x = torch.randn(100).cuda()
   y = torch.randn(100).cuda()
   # 删除张量引用，但CUDA内存并没有被释放
   x = None
   y = None
   # 使用empty_cache()释放未使用的CUDA内存
   torch.cuda.empty_cache()

3. 关闭PyTorch上下文。如果您使用的是PyTorch的分布式训练或分布式推理环境，那么每个进程可能会有多个上下文，您可以在不再需要时关闭它们。这对于释放由PyTorch管理的CUDA内存非常有用。如果您在会话期间不关闭上下文，则可能发生在使用大量GPU内存的情况下，会话之间的GPU内存共享。这可能会导致一个会话中的问题影响其他会话。在关闭上下文之前，请确保您不再需要任何使用该上下文的张量。一旦上下文被关闭，你将无法再访问那些张量。关闭上下文的一种方法是调用 torch.cuda.close()。然而这可能不是一个理想的解决方案，因为您需要管理上下文的生命周期并确保它们在不再需要时被关闭。一个更简单的方法是使用Python的 with 语句和 torch.cuda.device() 上下文管理器。这将自动关闭上下文并在必要时引发异常。例如：

   import torch
   # 在GPU上创建一个设备上下文并自动关闭它
   with torch.cuda.device(0):
   x = torch.randn(100).cuda()

   在以上代码中，torch.cuda.device(0) 创建了一个新的设备上下文并使其成为当前的设备上下文。一旦退出 with 块，该上下文就会被自动关闭。如果在上下文关闭后尝试访问其中的张量，将会引发一个异常。