使用Torch与Numba清理GPU内存后复用GPU报错的解决方法

核心问题 #

你用Numba的cuda.device.reset()或cuda.close()强制清理GPU内存的操作，会破坏PyTorch正在使用的CUDA上下文。CUDA上下文是进程级别的核心资源，PyTorch启动时会创建专属的CUDA上下文，一旦被Numba强制重置/关闭，后续PyTorch的GPU操作（比如model.to(device)）会因上下文失效抛出invalid argument错误。

正确的GPU内存清理方案 #

不要混用Numba和PyTorch的CUDA设备管理，优先使用PyTorch原生工具解决内存问题：

1. 标准内存清理流程 #

先确保所有不再需要的GPU张量被显式释放，再执行垃圾回收和缓存清理：

import gc
import torch
from torchvision import models

# 1. 显式删除不再使用的模型/张量
del model  # 释放之前的模型引用
# 若存在其他GPU张量，需一并用del删除

# 2. 触发Python垃圾回收，释放张量的CPU引用
gc.collect()

# 3. 清理PyTorch持有的GPU缓存
torch.cuda.empty_cache()

# 之后正常加载新模型
model = models.inception_v3(pretrained=True)
model.to(device)

2. 排查内存泄漏的根本原因 #

如果上述操作仍无法解决内存泄漏，需定位泄漏点：

• 检查是否有未被正确释放的中间张量：比如推理输出、损失张量未被del或覆盖，或被全局变量/闭包持有引用。
• 检查DataLoader的pin_memory参数：开启pin_memory=True后，需确保数据加载完成后无残留锁页内存。
• 检查模型的钩子函数（hook）：自定义的forward/backward钩子可能持有张量引用，导致无法被回收。
• 使用torch.cuda.memory_summary()查看内存使用详情，定位异常占用的模块。

3. 若必须混用Numba与PyTorch #

如果业务场景需要同时使用两者，需遵循以下规则：

• 在PyTorch初始化CUDA上下文之前完成Numba的GPU操作，再启动PyTorch。
• 若必须在PyTorch之后使用Numba并重置设备，重置后需重新初始化PyTorch的CUDA上下文：

  from numba import cuda
  import torch
  
  # 重置Numba的CUDA设备
  nb_device = cuda.get_current_device()
  nb_device.reset()
  
  # 重新初始化PyTorch的CUDA上下文
  torch.cuda.init()
  device = torch.device("cuda:0")
  
  # 之后正常加载模型
  model = models.inception_v3(pretrained=True)
  model.to(device)
  

错误方案的本质 #

你尝试的Numba清理操作直接操作CUDA设备底层资源，完全绕过了PyTorch的内存管理机制，导致PyTorch的CUDA上下文彻底失效，后续所有GPU操作必然报错。

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