PyTorch使用模型后如何彻底释放GPU显存？

你操作后仍有显存残留的核心原因：只删除了模型对象，训练过程中产生的其他CUDA侧对象（优化器、中间张量、pinned内存等）仍被Python引用持有，导致PyTorch无法释放对应显存，同时CUDA驱动会为运行中的Python进程预留部分显存上下文，不会完全退回到进程启动初始状态。

1. 全量清理CUDA关联对象 #

• 除model外，需要同步删除所有移动到CUDA侧的对象：优化器、损失函数对象、学习率调度器、训练/验证过程中保存的输出张量、标签张量、自定义CUDA缓存变量等
• 清理时注意执行顺序：先做Python垃圾回收清除无效引用，再清空CUDA缓存，示例代码如下：

import gc
import torch

# 替换为你自己定义的所有CUDA侧对象
del model, optimizer, criterion, lr_scheduler, train_outputs, val_labels
del train_dataloader, val_dataloader

gc.collect()
torch.cuda.empty_cache()

2. 特殊场景适配 #

• Jupyter环境：不要直接输出CUDA张量作为cell运行结果，Jupyter会自动持有该对象的引用导致无法释放，需要查看的话先调用.cpu()转为CPU张量再输出
• 多卡训练场景：清理前先销毁分布式进程组，添加代码：torch.distributed.destroy_process_group()
• DataLoader开启pin_memory：训练结束后主动删除dataloader和对应的dataset对象，解除pinned内存占用

3. 进程级显存强制清理 #

如果上述操作后仍有显存残留，且不想终止当前Python进程，可以调用PyTorch内置的CUDA设备重置接口：

# 重置指定CUDA设备的所有状态，释放所有该进程占用的显存
torch.cuda.device("cuda:0").reset()

注意：执行设备重置后，当前进程中所有已经创建的CUDA张量都会失效，不可再访问。

如果你的Python进程还在运行，CUDA驱动会预留100-300MiB左右的上下文显存，这部分是正常预留，只要后续新的训练任务可以正常占用显存就无需额外处理，只有当你需要把显存让给其他独立进程使用时才需要完全释放。

内容的提问来源于stack exchange，提问作者user3668129