在 PyTorch 中，训练完模型后，如果想要释放 GPU 内存，你可以使用以下几种方法：

1. 清空缓存

使用 torch.cuda.empty_cache() 来释放未使用的显存。请注意，这不会删除变量，只是帮助 PyTorch 释放未使用的内存。示例如下：

import torch  

# 在训练完成后  
torch.cuda.empty_cache()
  • 1.
  • 2.
  • 3.
  • 4.

2. 删除模型和变量

确保你删除不再使用的模型、优化器或其他变量，并调用 Python 的垃圾回收：

import gc  

# 假设你有一个模型和优化器  
del model  
del optimizer  
gc.collect()  # 垃圾回收
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.

3. 使用 with torch.no_grad()

在测试或推理模型时，使用 with torch.no_grad(): 来减少内存使用：

with torch.no_grad():  
    # 进行推理  
    output = model(input)
  • 1.
  • 2.
  • 3.

4. 切换到 CPU

如果你需要完全释放 GPU 内存，考虑将模型移到 CPU：

model.to('cpu')
  • 1.

5. 重启 Jupyter Notebook 或 Python 会话

如果显存仍然没有被有效释放，最直接的方法是重启你的 Python 会话或 Jupyter Notebook。

总结

  • 使用 torch.cuda.empty_cache() 来释放未使用的显存。
  • 及时删除不再使用的模型和变量。
  • 在评估模型时，使用 with torch.no_grad() 来减少内存占用。
  • 如果仍然面临内存问题，考虑重启整个会话