【Bug已解决】How to clear GPU memory after PyTorch model training without restarting kernel 解决方案问题描述在 PyTorch 深度学习开发中GPU 显存泄漏和占用不释放是最常见的工程问题之一。典型场景包括训练完成后 GPU 显存不释放模型训练结束后GPU 显存仍然被占用无法用于其他任务。Jupyter Notebook 中反复训练导致 OOM在 Notebook 中多次运行训练代码每次都占用更多显存最终导致CUDA out of memory。模型推理后显存不释放加载模型进行推理后显存没有被正确释放。异常中断后显存泄漏训练过程中出现错误中断GPU 显存没有被释放。这些问题的根本原因在于 Python 的垃圾回收机制和 PyTorch 的 CUDA 缓存策略。本文将系统地介绍如何正确释放 GPU 显存。错误复现错误一训练后 GPU 显存不释放import torch import torch.nn as nn # 训练模型 model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ).cuda() optimizer torch.optim.Adam(model.parameters()) criterion nn.CrossEntropyLoss() # 模拟训练 for epoch in range(10): x torch.randn(256, 1000).cuda() y torch.randint(0, 10, (256,)).cuda() optimizer.zero_grad() loss criterion(model(x), y) loss.backward() optimizer.step() print(Training done.) # 检查 GPU 显存 print(fAllocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB) print(fCached: {torch.cuda.memory_reserved() / 1024**2:.2f} MB) # 问题即使训练结束了GPU 显存仍然被占用 # del model # 忘记删除模型 # optimizer 还引用着模型的参数错误二Jupyter Notebook 中反复运行导致 OOM# 在 Jupyter Notebook 的不同 cell 中反复运行 # Cell 1: 第一次训练 model1 LargeModel().cuda() train(model1, ...) # GPU 显存: 4000 MB # Cell 2: 第二次训练没有释放 model1 model2 LargeModel().cuda() # 又分配了 4000 MB train(model2, ...) # GPU 显存: 8000 MB # Cell 3: 第三次训练 model3 LargeModel().cuda() # CUDA out of memory!错误三变量引用导致无法释放import torch import torch.nn as nn # 创建模型和数据 model LargeModel().cuda() optimizer torch.optim.Adam(model.parameters()) # 训练... losses [] for epoch in range(10): x torch.randn(256, 1000).cuda() output model(x) loss criterion(output, y) # 错误将 GPU 张量保存在列表中 losses.append(loss) # loss 是 GPU 张量不会被垃圾回收 loss.backward() optimizer.step() # losses 列表持有所有 loss 张量的引用 # 这些张量占用的 GPU 显存不会被释放根因分析1. PyTorch CUDA 内存管理机制PyTorch 使用一个内存分配器CUDACachingAllocator来管理 GPU 显存已分配Allocated当前被张量实际使用的显存。缓存Cached/ReservedPyTorch 预留的显存用于未来的分配请求。当张量被删除时PyTorch 不会立即将显存归还给 GPU 驱动而是将其放入缓存池以便后续快速复用。这就是为什么torch.cuda.memory_allocated()和torch.cuda.memory_reserved()的值不同的原因。2. Python 垃圾回收与引用计数Python 使用引用计数来管理对象的生命周期。一个对象只有在所有引用都被删除后才会被垃圾回收model Model().cuda() # 引用计数 1 another_ref model # 引用计数 2 del model # 引用计数 1对象仍然存在 del another_ref # 引用计数 0对象被回收常见的隐藏引用优化器optimizer持有模型参数的引用列表/字典中保存的张量计算图中的中间变量全局变量Jupyter Notebook 的输出历史3. 计算图未释放如果张量的requires_gradTrue且计算图没有被释放中间变量会一直保存在内存中x torch.randn(100, 100, requires_gradTrue).cuda() y model(x) # y 保存了整个计算图包括所有中间激活值 # 只有调用 y.backward() 或 del y 后才会释放4. Jupyter Notebook 的特殊问题Jupyter Notebook 会保存所有 cell 的输出和变量历史。即使你del了一个变量Notebook 的输出缓存_,__,_oh等可能仍然持有引用。解决方案方案一正确删除变量并清空缓存import torch import torch.nn as nn import gc def release_gpu_memory(*variables): 释放 GPU 显存。 Args: *variables: 需要删除的变量名 # 1. 删除变量 for var in variables: del var # 2. 触发垃圾回收 gc.collect() # 3. 清空 PyTorch 缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() # 打印显存信息 allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(fAfter cleanup - Allocated: {allocated:.2f} MB, Cached: {cached:.2f} MB) # 使用示例 model nn.Linear(1000, 10).cuda() optimizer torch.optim.Adam(model.parameters()) # 训练... # 释放显存 release_gpu_memory(model, optimizer)方案二使用上下文管理器import torch import gc from contextlib import contextmanager contextmanager def gpu_memory_manager(devicecuda, verboseTrue): GPU 显存管理上下文管理器。 使用方式: with gpu_memory_manager(): model LargeModel().cuda() train(model) # 退出 with 块后自动释放显存 if verbose and torch.cuda.is_available(): allocated_before torch.cuda.memory_allocated() / 1024**2 cached_before torch.cuda.memory_reserved() / 1024**2 print(fBefore: Allocated{allocated_before:.2f}MB, Cached{cached_before:.2f}MB) # 记录进入前的变量 initial_vars set(dir()) try: yield finally: # 清理新创建的变量 current_vars set(dir()) new_vars current_vars - initial_vars for var_name in new_vars: obj eval(var_name) if isinstance(obj, torch.Tensor) and obj.is_cuda: del obj # 垃圾回收 gc.collect() # 清空缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() if verbose: allocated_after torch.cuda.memory_allocated() / 1024**2 cached_after torch.cuda.memory_reserved() / 1024**2 print(fAfter: Allocated{allocated_after:.2f}MB, Cached{cached_after:.2f}MB) # 使用示例 with gpu_memory_manager(): model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 10) ).cuda() x torch.randn(256, 1000).cuda() output model(x) print(fOutput shape: {output.shape}) # 退出 with 块后自动清理方案三完整的训练-释放流程import torch import torch.nn as nn import gc def train_and_release(): 完整的训练-释放流程。 所有变量都在函数内部定义函数结束后自动释放。 # 所有变量都是局部变量 model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ).cuda() optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion nn.CrossEntropyLoss() # 训练 model.train() for epoch in range(5): x torch.randn(256, 1000).cuda() y torch.randint(0, 10, (256,)).cuda() optimizer.zero_grad() output model(x) loss criterion(output, y) loss.backward() optimizer.step() # 重要只保存标量值不保存 GPU 张量 loss_value loss.item() # .item() 返回 Python float print(fEpoch {epoch}: loss{loss_value:.4f}) # 手动删除不需要的张量 del output, loss # 保存模型移到 CPU 再保存 model_cpu model.cpu() torch.save(model_cpu.state_dict(), model.pth) # 清理 del model, optimizer, model_cpu gc.collect() torch.cuda.empty_cache() print(fFinal GPU memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB) # 调用函数 train_and_release() # 函数返回后所有局部变量自动被回收完整修复代码import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset import gc import sys # # 完整示例GPU 显存管理与释放 # class GPUMemoryTracker: GPU 显存跟踪器。 用于监控和调试 GPU 显存使用情况。 def __init__(self, devicecuda): self.device device self.snapshots [] def snapshot(self, label): 记录当前显存使用情况。 if not torch.cuda.is_available(): return allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 max_allocated torch.cuda.max_memory_allocated() / 1024**2 snapshot { label: label, allocated_mb: allocated, cached_mb: cached, max_allocated_mb: max_allocated } self.snapshots.append(snapshot) print(f[{label}] Allocated: {allocated:.2f} MB, fCached: {cached:.2f} MB, fPeak: {max_allocated:.2f} MB) return snapshot def reset_peak(self): 重置峰值统计。 if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def print_summary(self): 打印所有快照的摘要。 print(\n * 60) print(GPU Memory Summary) print( * 60) for snap in self.snapshots: print(f {snap[label]:30} fAlloc: {snap[allocated_mb]:8.2f} MB | fCache: {snap[cached_mb]:8.2f} MB | fPeak: {snap[max_allocated_mb]:8.2f} MB) print( * 60) def find_gpu_tensors(): 查找当前所有 GPU 张量。 用于调试显存泄漏问题。 gpu_tensors [] for obj in gc.get_objects(): try: if (isinstance(obj, torch.Tensor) and obj.is_cuda and id(obj) not in [id(t) for t, _ in gpu_tensors]): gpu_tensors.append((obj, obj.element_size() * obj.nelement())) except: pass total_bytes sum(size for _, size in gpu_tensors) print(fFound {len(gpu_tensors)} GPU tensors, total: {total_bytes / 1024**2:.2f} MB) # 按大小排序显示最大的几个 gpu_tensors.sort(keylambda x: x[1], reverseTrue) for tensor, size in gpu_tensors[:5]: print(f Shape: {tensor.shape}, Size: {size / 1024**2:.2f} MB) return gpu_tensors def full_cleanup(): 执行完整的 GPU 清理。 # 1. 触发 Python 垃圾回收 gc.collect() if torch.cuda.is_available(): # 2. 同步 CUDA 操作 torch.cuda.synchronize() # 3. 清空缓存 torch.cuda.empty_cache() # 4. 重置峰值统计 torch.cuda.reset_peak_memory_stats() # 5. IPC 收集如果使用了多进程 if hasattr(torch.cuda, ipc_collect): torch.cuda.ipc_collect() allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(fCleanup complete - Allocated: {allocated:.2f} MB, Cached: {cached:.2f} MB) class ModelTrainer: 完整的模型训练器包含正确的显存管理。 def __init__(self, model, devicecuda): self.model model self.device device self.optimizer None self.criterion nn.CrossEntropyLoss() self.tracker GPUMemoryTracker(device) def train(self, train_loader, num_epochs10, lr0.001): 训练模型。 self.model self.model.to(self.device) self.optimizer torch.optim.AdamW( self.model.parameters(), lrlr, weight_decay0.01 ) self.tracker.reset_peak() self.tracker.snapshot(Training start) for epoch in range(num_epochs): self.model.train() epoch_loss 0.0 num_batches 0 for batch_idx, (inputs, targets) in enumerate(train_loader): inputs inputs.to(self.device) targets targets.to(self.device) self.optimizer.zero_grad() outputs self.model(inputs) loss self.criterion(outputs, targets) loss.backward() # 梯度裁剪 torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0) self.optimizer.step() # 只保存标量值 epoch_loss loss.item() num_batches 1 # 显式删除中间变量 del outputs, loss avg_loss epoch_loss / num_batches if (epoch 1) % 5 0: self.tracker.snapshot(fEpoch {epoch1}) print(f Loss: {avg_loss:.4f}) self.tracker.snapshot(Training end) def evaluate(self, test_loader): 评估模型。 self.model.eval() correct 0 total 0 self.tracker.snapshot(Eval start) with torch.no_grad(): # 不保存计算图 for inputs, targets in test_loader: inputs inputs.to(self.device) targets targets.to(self.device) outputs self.model(inputs) _, predicted outputs.max(1) total targets.size(0) correct (predicted targets).sum().item() del outputs, predicted accuracy correct / total self.tracker.snapshot(Eval end) print(fTest Accuracy: {accuracy:.4f}) return accuracy def save_model(self, filepath): 保存模型移到 CPU 再保存。 # 将模型移到 CPU 再保存避免加载时设备问题 self.model self.model.cpu() torch.save(self.model.state_dict(), filepath) print(fModel saved to {filepath}) # 移回 GPU如果需要继续训练 self.model self.model.to(self.device) def cleanup(self): 清理所有资源。 del self.optimizer del self.criterion self.model self.model.cpu() del self.model full_cleanup() self.tracker.snapshot(After cleanup) self.tracker.print_summary() def demo_jupyter_cleanup(): 演示 Jupyter Notebook 中的清理方法。 print(\n * 60) print(Jupyter Notebook 清理方法) print( * 60) # 方法 1使用 ipython 的变量清理 # 在 Jupyter cell 中运行 # # 清除所有变量 # %reset -f # # 或者选择性清除 # for name in dir(): # if not name.startswith(_): # del globals()[name] # 方法 2使用魔法命令 # # 在 cell 开头添加 # %xdel model # 删除变量并清除所有引用 # 方法 3完整的清理函数 def jupyter_full_cleanup(): Jupyter Notebook 完整清理。 import gc import torch # 清除 IPython 的输出缓存 try: from IPython import get_ipython ipython get_ipython() if ipython is not None: ipython.history_manager.reset() except: pass # 垃圾回收 gc.collect() # 清空 CUDA 缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() print(Jupyter cleanup complete.) jupyter_full_cleanup() def main(): 主函数。 if not torch.cuda.is_available(): print(CUDA not available. Running CPU demo.) device cpu else: print(fUsing GPU: {torch.cuda.get_device_name(0)}) device cuda torch.manual_seed(42) # 生成数据 num_samples 2000 X torch.randn(num_samples, 100) y torch.randint(0, 10, (num_samples,)) split int(0.8 * num_samples) train_dataset TensorDataset(X[:split], y[:split]) test_dataset TensorDataset(X[split:], y[split:]) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse) # 创建模型 model nn.Sequential( nn.Linear(100, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, 10) ) # 训练 trainer ModelTrainer(model, devicedevice) trainer.train(train_loader, num_epochs10, lr0.001) trainer.evaluate(test_loader) trainer.save_model(model.pth) # 清理 trainer.cleanup() # Jupyter 清理演示 demo_jupyter_cleanup() print(\nDone!) if __name__ __main__: main()运行输出示例Using GPU: NVIDIA GeForce RTX 3090 [Training start] Allocated: 0.25 MB, Cached: 2.00 MB, Peak: 0.25 MB [Epoch 5] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Loss: 1.8234 [Epoch 10] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Loss: 1.2345 [Training end] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB [Eval start] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB [Eval end] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Test Accuracy: 0.7250 Model saved to model.pth Cleanup complete - Allocated: 0.00 MB, Cached: 0.00 MB [After cleanup] Allocated: 0.00 MB, Cached: 0.00 MB, Peak: 0.00 MB GPU Memory Summary Training start Alloc: 0.25 MB | Cache: 2.00 MB | Peak: 0.25 MB Epoch 5 Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Epoch 10 Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Training end Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Eval start Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Eval end Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB After cleanup Alloc: 0.00 MB | Cache: 0.00 MB | Peak: 0.00 MB 常见陷阱与注意事项陷阱 1只调用empty_cache不删除变量# 错误只清空缓存不删除变量 torch.cuda.empty_cache() # 缓存被清空了但已分配的显存仍然被占用 # 正确先删除变量再清空缓存 del model, optimizer gc.collect() torch.cuda.empty_cache()陷阱 2保存 GPU 张量到列表# 错误保存 GPU 张量 losses [] for batch in dataloader: loss criterion(model(batch), target) losses.append(loss) # GPU 张量不会被回收 # 正确只保存标量值 losses [] for batch in dataloader: loss criterion(model(batch), target) losses.append(loss.item()) # Python float不占 GPU 显存陷阱 3忘记torch.no_grad()在推理时# 错误推理时不使用 no_grad model.eval() output model(x) # 保存了计算图占用额外显存 # 正确推理时使用 no_grad model.eval() with torch.no_grad(): output model(x) # 不保存计算图陷阱 4optimizer 持有模型引用# 错误只删除模型不删除优化器 model Model().cuda() optimizer torch.optim.Adam(model.parameters()) # ... del model # optimizer 仍然引用着模型的参数 gc.collect() torch.cuda.empty_cache() # 显存可能不会被完全释放 # 正确同时删除优化器 del model, optimizer gc.collect() torch.cuda.empty_cache()陷阱 5del不等于立即释放# del 只是减少引用计数 # 如果还有其他引用对象不会被回收 model Model().cuda() ref model # 另一个引用 del model # 引用计数从 2 变为 1对象仍然存在 gc.collect() torch.cuda.empty_cache() # 显存仍然被占用 del ref # 引用计数变为 0对象被回收 gc.collect() torch.cuda.empty_cache() # 现在显存才被释放陷阱 6torch.cuda.empty_cache()的实际效果# empty_cache 只释放缓存reserved不释放已分配allocated # 它将 PyTorch 缓存池中未使用的显存归还给 GPU 驱动 # 如果 allocated 1000 MB, cached 1500 MB # 调用 empty_cache 后allocated 1000 MB, cached 1000 MB # 释放了 500 MB 的缓存但 1000 MB 的已分配显存不变 # 要释放 allocated必须先删除引用该显存的张量总结正确释放 GPU 显存需要遵循以下核心步骤删除所有引用del model, optimizer, loss, tensor等确保引用计数归零。触发垃圾回收gc.collect()强制 Python 回收不可达对象。清空 CUDA 缓存torch.cuda.empty_cache()将未使用的缓存归还给 GPU 驱动。使用torch.no_grad()推理时不保存计算图减少显存占用。保存标量而非张量用.item()提取标量值避免持有 GPU 张量引用。函数封装将训练代码封装在函数中利用局部变量的自动回收机制。上下文管理器使用with语句自动管理资源的分配和释放。Jupyter Notebook 特殊处理注意 IPython 的输出缓存使用%reset或%xdel。记住torch.cuda.empty_cache()只是锦上添花真正释放显存的是删除变量引用。养成良好的显存管理习惯可以避免绝大多数 OOM 问题。 【Bug已解决】How to clear GPU memory after PyTorch model training without restarting kernel 解决方案
【Bug已解决】How to clear GPU memory after PyTorch model training without restarting kernel 解决方案问题描述在 PyTorch 深度学习开发中GPU 显存泄漏和占用不释放是最常见的工程问题之一。典型场景包括训练完成后 GPU 显存不释放模型训练结束后GPU 显存仍然被占用无法用于其他任务。Jupyter Notebook 中反复训练导致 OOM在 Notebook 中多次运行训练代码每次都占用更多显存最终导致CUDA out of memory。模型推理后显存不释放加载模型进行推理后显存没有被正确释放。异常中断后显存泄漏训练过程中出现错误中断GPU 显存没有被释放。这些问题的根本原因在于 Python 的垃圾回收机制和 PyTorch 的 CUDA 缓存策略。本文将系统地介绍如何正确释放 GPU 显存。错误复现错误一训练后 GPU 显存不释放import torch import torch.nn as nn # 训练模型 model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ).cuda() optimizer torch.optim.Adam(model.parameters()) criterion nn.CrossEntropyLoss() # 模拟训练 for epoch in range(10): x torch.randn(256, 1000).cuda() y torch.randint(0, 10, (256,)).cuda() optimizer.zero_grad() loss criterion(model(x), y) loss.backward() optimizer.step() print(Training done.) # 检查 GPU 显存 print(fAllocated: {torch.cuda.memory_allocated() / 1024**2:.2f} MB) print(fCached: {torch.cuda.memory_reserved() / 1024**2:.2f} MB) # 问题即使训练结束了GPU 显存仍然被占用 # del model # 忘记删除模型 # optimizer 还引用着模型的参数错误二Jupyter Notebook 中反复运行导致 OOM# 在 Jupyter Notebook 的不同 cell 中反复运行 # Cell 1: 第一次训练 model1 LargeModel().cuda() train(model1, ...) # GPU 显存: 4000 MB # Cell 2: 第二次训练没有释放 model1 model2 LargeModel().cuda() # 又分配了 4000 MB train(model2, ...) # GPU 显存: 8000 MB # Cell 3: 第三次训练 model3 LargeModel().cuda() # CUDA out of memory!错误三变量引用导致无法释放import torch import torch.nn as nn # 创建模型和数据 model LargeModel().cuda() optimizer torch.optim.Adam(model.parameters()) # 训练... losses [] for epoch in range(10): x torch.randn(256, 1000).cuda() output model(x) loss criterion(output, y) # 错误将 GPU 张量保存在列表中 losses.append(loss) # loss 是 GPU 张量不会被垃圾回收 loss.backward() optimizer.step() # losses 列表持有所有 loss 张量的引用 # 这些张量占用的 GPU 显存不会被释放根因分析1. PyTorch CUDA 内存管理机制PyTorch 使用一个内存分配器CUDACachingAllocator来管理 GPU 显存已分配Allocated当前被张量实际使用的显存。缓存Cached/ReservedPyTorch 预留的显存用于未来的分配请求。当张量被删除时PyTorch 不会立即将显存归还给 GPU 驱动而是将其放入缓存池以便后续快速复用。这就是为什么torch.cuda.memory_allocated()和torch.cuda.memory_reserved()的值不同的原因。2. Python 垃圾回收与引用计数Python 使用引用计数来管理对象的生命周期。一个对象只有在所有引用都被删除后才会被垃圾回收model Model().cuda() # 引用计数 1 another_ref model # 引用计数 2 del model # 引用计数 1对象仍然存在 del another_ref # 引用计数 0对象被回收常见的隐藏引用优化器optimizer持有模型参数的引用列表/字典中保存的张量计算图中的中间变量全局变量Jupyter Notebook 的输出历史3. 计算图未释放如果张量的requires_gradTrue且计算图没有被释放中间变量会一直保存在内存中x torch.randn(100, 100, requires_gradTrue).cuda() y model(x) # y 保存了整个计算图包括所有中间激活值 # 只有调用 y.backward() 或 del y 后才会释放4. Jupyter Notebook 的特殊问题Jupyter Notebook 会保存所有 cell 的输出和变量历史。即使你del了一个变量Notebook 的输出缓存_,__,_oh等可能仍然持有引用。解决方案方案一正确删除变量并清空缓存import torch import torch.nn as nn import gc def release_gpu_memory(*variables): 释放 GPU 显存。 Args: *variables: 需要删除的变量名 # 1. 删除变量 for var in variables: del var # 2. 触发垃圾回收 gc.collect() # 3. 清空 PyTorch 缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() # 打印显存信息 allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(fAfter cleanup - Allocated: {allocated:.2f} MB, Cached: {cached:.2f} MB) # 使用示例 model nn.Linear(1000, 10).cuda() optimizer torch.optim.Adam(model.parameters()) # 训练... # 释放显存 release_gpu_memory(model, optimizer)方案二使用上下文管理器import torch import gc from contextlib import contextmanager contextmanager def gpu_memory_manager(devicecuda, verboseTrue): GPU 显存管理上下文管理器。 使用方式: with gpu_memory_manager(): model LargeModel().cuda() train(model) # 退出 with 块后自动释放显存 if verbose and torch.cuda.is_available(): allocated_before torch.cuda.memory_allocated() / 1024**2 cached_before torch.cuda.memory_reserved() / 1024**2 print(fBefore: Allocated{allocated_before:.2f}MB, Cached{cached_before:.2f}MB) # 记录进入前的变量 initial_vars set(dir()) try: yield finally: # 清理新创建的变量 current_vars set(dir()) new_vars current_vars - initial_vars for var_name in new_vars: obj eval(var_name) if isinstance(obj, torch.Tensor) and obj.is_cuda: del obj # 垃圾回收 gc.collect() # 清空缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() if verbose: allocated_after torch.cuda.memory_allocated() / 1024**2 cached_after torch.cuda.memory_reserved() / 1024**2 print(fAfter: Allocated{allocated_after:.2f}MB, Cached{cached_after:.2f}MB) # 使用示例 with gpu_memory_manager(): model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 10) ).cuda() x torch.randn(256, 1000).cuda() output model(x) print(fOutput shape: {output.shape}) # 退出 with 块后自动清理方案三完整的训练-释放流程import torch import torch.nn as nn import gc def train_and_release(): 完整的训练-释放流程。 所有变量都在函数内部定义函数结束后自动释放。 # 所有变量都是局部变量 model nn.Sequential( nn.Linear(1000, 512), nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 10) ).cuda() optimizer torch.optim.Adam(model.parameters(), lr0.001) criterion nn.CrossEntropyLoss() # 训练 model.train() for epoch in range(5): x torch.randn(256, 1000).cuda() y torch.randint(0, 10, (256,)).cuda() optimizer.zero_grad() output model(x) loss criterion(output, y) loss.backward() optimizer.step() # 重要只保存标量值不保存 GPU 张量 loss_value loss.item() # .item() 返回 Python float print(fEpoch {epoch}: loss{loss_value:.4f}) # 手动删除不需要的张量 del output, loss # 保存模型移到 CPU 再保存 model_cpu model.cpu() torch.save(model_cpu.state_dict(), model.pth) # 清理 del model, optimizer, model_cpu gc.collect() torch.cuda.empty_cache() print(fFinal GPU memory: {torch.cuda.memory_allocated() / 1024**2:.2f} MB) # 调用函数 train_and_release() # 函数返回后所有局部变量自动被回收完整修复代码import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset import gc import sys # # 完整示例GPU 显存管理与释放 # class GPUMemoryTracker: GPU 显存跟踪器。 用于监控和调试 GPU 显存使用情况。 def __init__(self, devicecuda): self.device device self.snapshots [] def snapshot(self, label): 记录当前显存使用情况。 if not torch.cuda.is_available(): return allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 max_allocated torch.cuda.max_memory_allocated() / 1024**2 snapshot { label: label, allocated_mb: allocated, cached_mb: cached, max_allocated_mb: max_allocated } self.snapshots.append(snapshot) print(f[{label}] Allocated: {allocated:.2f} MB, fCached: {cached:.2f} MB, fPeak: {max_allocated:.2f} MB) return snapshot def reset_peak(self): 重置峰值统计。 if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def print_summary(self): 打印所有快照的摘要。 print(\n * 60) print(GPU Memory Summary) print( * 60) for snap in self.snapshots: print(f {snap[label]:30} fAlloc: {snap[allocated_mb]:8.2f} MB | fCache: {snap[cached_mb]:8.2f} MB | fPeak: {snap[max_allocated_mb]:8.2f} MB) print( * 60) def find_gpu_tensors(): 查找当前所有 GPU 张量。 用于调试显存泄漏问题。 gpu_tensors [] for obj in gc.get_objects(): try: if (isinstance(obj, torch.Tensor) and obj.is_cuda and id(obj) not in [id(t) for t, _ in gpu_tensors]): gpu_tensors.append((obj, obj.element_size() * obj.nelement())) except: pass total_bytes sum(size for _, size in gpu_tensors) print(fFound {len(gpu_tensors)} GPU tensors, total: {total_bytes / 1024**2:.2f} MB) # 按大小排序显示最大的几个 gpu_tensors.sort(keylambda x: x[1], reverseTrue) for tensor, size in gpu_tensors[:5]: print(f Shape: {tensor.shape}, Size: {size / 1024**2:.2f} MB) return gpu_tensors def full_cleanup(): 执行完整的 GPU 清理。 # 1. 触发 Python 垃圾回收 gc.collect() if torch.cuda.is_available(): # 2. 同步 CUDA 操作 torch.cuda.synchronize() # 3. 清空缓存 torch.cuda.empty_cache() # 4. 重置峰值统计 torch.cuda.reset_peak_memory_stats() # 5. IPC 收集如果使用了多进程 if hasattr(torch.cuda, ipc_collect): torch.cuda.ipc_collect() allocated torch.cuda.memory_allocated() / 1024**2 cached torch.cuda.memory_reserved() / 1024**2 print(fCleanup complete - Allocated: {allocated:.2f} MB, Cached: {cached:.2f} MB) class ModelTrainer: 完整的模型训练器包含正确的显存管理。 def __init__(self, model, devicecuda): self.model model self.device device self.optimizer None self.criterion nn.CrossEntropyLoss() self.tracker GPUMemoryTracker(device) def train(self, train_loader, num_epochs10, lr0.001): 训练模型。 self.model self.model.to(self.device) self.optimizer torch.optim.AdamW( self.model.parameters(), lrlr, weight_decay0.01 ) self.tracker.reset_peak() self.tracker.snapshot(Training start) for epoch in range(num_epochs): self.model.train() epoch_loss 0.0 num_batches 0 for batch_idx, (inputs, targets) in enumerate(train_loader): inputs inputs.to(self.device) targets targets.to(self.device) self.optimizer.zero_grad() outputs self.model(inputs) loss self.criterion(outputs, targets) loss.backward() # 梯度裁剪 torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5.0) self.optimizer.step() # 只保存标量值 epoch_loss loss.item() num_batches 1 # 显式删除中间变量 del outputs, loss avg_loss epoch_loss / num_batches if (epoch 1) % 5 0: self.tracker.snapshot(fEpoch {epoch1}) print(f Loss: {avg_loss:.4f}) self.tracker.snapshot(Training end) def evaluate(self, test_loader): 评估模型。 self.model.eval() correct 0 total 0 self.tracker.snapshot(Eval start) with torch.no_grad(): # 不保存计算图 for inputs, targets in test_loader: inputs inputs.to(self.device) targets targets.to(self.device) outputs self.model(inputs) _, predicted outputs.max(1) total targets.size(0) correct (predicted targets).sum().item() del outputs, predicted accuracy correct / total self.tracker.snapshot(Eval end) print(fTest Accuracy: {accuracy:.4f}) return accuracy def save_model(self, filepath): 保存模型移到 CPU 再保存。 # 将模型移到 CPU 再保存避免加载时设备问题 self.model self.model.cpu() torch.save(self.model.state_dict(), filepath) print(fModel saved to {filepath}) # 移回 GPU如果需要继续训练 self.model self.model.to(self.device) def cleanup(self): 清理所有资源。 del self.optimizer del self.criterion self.model self.model.cpu() del self.model full_cleanup() self.tracker.snapshot(After cleanup) self.tracker.print_summary() def demo_jupyter_cleanup(): 演示 Jupyter Notebook 中的清理方法。 print(\n * 60) print(Jupyter Notebook 清理方法) print( * 60) # 方法 1使用 ipython 的变量清理 # 在 Jupyter cell 中运行 # # 清除所有变量 # %reset -f # # 或者选择性清除 # for name in dir(): # if not name.startswith(_): # del globals()[name] # 方法 2使用魔法命令 # # 在 cell 开头添加 # %xdel model # 删除变量并清除所有引用 # 方法 3完整的清理函数 def jupyter_full_cleanup(): Jupyter Notebook 完整清理。 import gc import torch # 清除 IPython 的输出缓存 try: from IPython import get_ipython ipython get_ipython() if ipython is not None: ipython.history_manager.reset() except: pass # 垃圾回收 gc.collect() # 清空 CUDA 缓存 if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() print(Jupyter cleanup complete.) jupyter_full_cleanup() def main(): 主函数。 if not torch.cuda.is_available(): print(CUDA not available. Running CPU demo.) device cpu else: print(fUsing GPU: {torch.cuda.get_device_name(0)}) device cuda torch.manual_seed(42) # 生成数据 num_samples 2000 X torch.randn(num_samples, 100) y torch.randint(0, 10, (num_samples,)) split int(0.8 * num_samples) train_dataset TensorDataset(X[:split], y[:split]) test_dataset TensorDataset(X[split:], y[split:]) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse) # 创建模型 model nn.Sequential( nn.Linear(100, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, 10) ) # 训练 trainer ModelTrainer(model, devicedevice) trainer.train(train_loader, num_epochs10, lr0.001) trainer.evaluate(test_loader) trainer.save_model(model.pth) # 清理 trainer.cleanup() # Jupyter 清理演示 demo_jupyter_cleanup() print(\nDone!) if __name__ __main__: main()运行输出示例Using GPU: NVIDIA GeForce RTX 3090 [Training start] Allocated: 0.25 MB, Cached: 2.00 MB, Peak: 0.25 MB [Epoch 5] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Loss: 1.8234 [Epoch 10] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Loss: 1.2345 [Training end] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB [Eval start] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB [Eval end] Allocated: 12.34 MB, Cached: 20.00 MB, Peak: 15.67 MB Test Accuracy: 0.7250 Model saved to model.pth Cleanup complete - Allocated: 0.00 MB, Cached: 0.00 MB [After cleanup] Allocated: 0.00 MB, Cached: 0.00 MB, Peak: 0.00 MB GPU Memory Summary Training start Alloc: 0.25 MB | Cache: 2.00 MB | Peak: 0.25 MB Epoch 5 Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Epoch 10 Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Training end Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Eval start Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB Eval end Alloc: 12.34 MB | Cache: 20.00 MB | Peak: 15.67 MB After cleanup Alloc: 0.00 MB | Cache: 0.00 MB | Peak: 0.00 MB 常见陷阱与注意事项陷阱 1只调用empty_cache不删除变量# 错误只清空缓存不删除变量 torch.cuda.empty_cache() # 缓存被清空了但已分配的显存仍然被占用 # 正确先删除变量再清空缓存 del model, optimizer gc.collect() torch.cuda.empty_cache()陷阱 2保存 GPU 张量到列表# 错误保存 GPU 张量 losses [] for batch in dataloader: loss criterion(model(batch), target) losses.append(loss) # GPU 张量不会被回收 # 正确只保存标量值 losses [] for batch in dataloader: loss criterion(model(batch), target) losses.append(loss.item()) # Python float不占 GPU 显存陷阱 3忘记torch.no_grad()在推理时# 错误推理时不使用 no_grad model.eval() output model(x) # 保存了计算图占用额外显存 # 正确推理时使用 no_grad model.eval() with torch.no_grad(): output model(x) # 不保存计算图陷阱 4optimizer 持有模型引用# 错误只删除模型不删除优化器 model Model().cuda() optimizer torch.optim.Adam(model.parameters()) # ... del model # optimizer 仍然引用着模型的参数 gc.collect() torch.cuda.empty_cache() # 显存可能不会被完全释放 # 正确同时删除优化器 del model, optimizer gc.collect() torch.cuda.empty_cache()陷阱 5del不等于立即释放# del 只是减少引用计数 # 如果还有其他引用对象不会被回收 model Model().cuda() ref model # 另一个引用 del model # 引用计数从 2 变为 1对象仍然存在 gc.collect() torch.cuda.empty_cache() # 显存仍然被占用 del ref # 引用计数变为 0对象被回收 gc.collect() torch.cuda.empty_cache() # 现在显存才被释放陷阱 6torch.cuda.empty_cache()的实际效果# empty_cache 只释放缓存reserved不释放已分配allocated # 它将 PyTorch 缓存池中未使用的显存归还给 GPU 驱动 # 如果 allocated 1000 MB, cached 1500 MB # 调用 empty_cache 后allocated 1000 MB, cached 1000 MB # 释放了 500 MB 的缓存但 1000 MB 的已分配显存不变 # 要释放 allocated必须先删除引用该显存的张量总结正确释放 GPU 显存需要遵循以下核心步骤删除所有引用del model, optimizer, loss, tensor等确保引用计数归零。触发垃圾回收gc.collect()强制 Python 回收不可达对象。清空 CUDA 缓存torch.cuda.empty_cache()将未使用的缓存归还给 GPU 驱动。使用torch.no_grad()推理时不保存计算图减少显存占用。保存标量而非张量用.item()提取标量值避免持有 GPU 张量引用。函数封装将训练代码封装在函数中利用局部变量的自动回收机制。上下文管理器使用with语句自动管理资源的分配和释放。Jupyter Notebook 特殊处理注意 IPython 的输出缓存使用%reset或%xdel。记住torch.cuda.empty_cache()只是锦上添花真正释放显存的是删除变量引用。养成良好的显存管理习惯可以避免绝大多数 OOM 问题。