如何通过Python函数删除全局变量以释放PyTorch GPU显存？

问题原因 #

当前的实现不生效是由Python的参数传递机制导致的：当你把变量传入函数时，实际上传递的是对象的引用，函数内部del arg仅删除了当前函数局部作用域内的引用变量，外层作用域（全局/函数外层）的原始变量仍然持有对象引用，对象引用计数未归零，因此不会被垃圾回收，显存也无法释放。

优化后的实现方案 #

我们可以改成传入变量的字符串名称，同时传入对应作用域的命名空间字典（locals() 或者 globals()），函数内部直接操作外层作用域的变量：

import torch
import gc

def free_gpu_memory(var_names: list, namespace: dict) -> None:
    """删除指定变量并释放GPU显存
    Args:
        var_names: 要删除的变量名列表，传字符串格式
        namespace: 变量所在的命名空间，局部变量传locals()，全局变量传globals()
    """
    for name in var_names:
        if name in namespace:
            del namespace[name]
    gc.collect()
    torch.cuda.empty_cache()

使用示例 #

局部变量场景（训练循环场景） #

def train_one_epoch():
    a = torch.tensor([1,2,3]).cuda()
    b = torch.tensor([4,5,6]).cuda()
    # 用完变量后调用，局部变量传当前作用域的locals()
    free_gpu_memory(["a", "b"], locals())

def train_loop(num: int):
    for i in range(num):
        train_one_epoch()

全局变量场景 #

# 全局作用域定义的变量
global_model = torch.nn.Linear(10, 2).cuda()
global_optim = torch.optim.Adam(global_model.parameters())

# 要删除全局变量时传globals()
free_gpu_memory(["global_model", "global_optim"], globals())

更推荐的简化写法 #

如果不想额外处理命名空间，也可以用更直观的逻辑，避开作用域问题：

# 清理逻辑只保留GC和显存清空操作
def free_gpu_memory() -> None:
    gc.collect()
    torch.cuda.empty_cache()

# 使用时先手动把需要释放的变量设为None，再调用清理即可
a = torch.tensor([1,2,3]).cuda()
b = torch.tensor([4,5,6]).cuda()
a = b = None
free_gpu_memory()

这种写法没有隐式的作用域操作，不易出现隐藏bug，更适合日常使用。

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