PyTorch循环迭代更新张量引发M1/CUDA GPU内存泄漏求助

你遇到的内存泄漏核心原因是：循环中持续创建新张量，旧张量的GPU内存未被及时回收——PyTorch的MPS/CUDA内存池不会主动释放未使用的空闲内存，加上递归赋值p1 = leak_example(p1, device)会让旧张量的引用在循环周期内暂时无法被Python GC处理，积累后导致内存持续增长。

以下是具体解决措施：

• 手动触发垃圾回收与设备内存清理
  在循环中定期调用Python垃圾回收和PyTorch设备内存清理函数，强制释放未被引用的GPU内存。empty_cache()会回收内存池中的空闲部分，配合GC能更高效释放内存。

• 移除不必要的detach操作
  你已经设置torch.set_grad_enabled(False)，所有张量默认无需计算梯度，detach()完全多余，只会额外创建张量副本增加内存负担。

• 原地更新张量，减少新对象创建
  避免每次返回新张量赋值给p1，改为在函数内部原地修改传入的张量，减少不必要的内存分配。若必须返回新张量，可显式删除旧张量后再赋值。

• 优化切片赋值逻辑
  a1[~mask] = torch.rand_like(a1)[~mask]会先创建完整随机张量再切片，产生冗余临时张量。可直接创建与~mask匹配的随机片段再赋值，降低内存占用。

修改后的示例代码 #

import torch
import gc

def leak_example(p1, device):
    # 移除不必要的detach操作
    t1 = torch.rand_like(p1, device=device)
    u1 = p1 + 2 * t1
    
    B = torch.rand_like(u1, device=device)
    mask = u1 < B
    
    # 原地更新p1，避免创建新张量
    p1.copy_(u1)
    # 优化赋值：仅生成需要的随机张量片段
    rand_vals = torch.rand(torch.sum(~mask), device=device)
    p1[~mask] = rand_vals
    return p1

if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
torch.set_grad_enabled(False)
        
p1 = torch.rand(5, 5, 224, 224, device=device)
for i in range(10000):
    p1 = leak_example(p1, device)
    # 每100次循环清理一次，频率可按需调整
    if i % 100 == 0:
        gc.collect()
        if device.type == "mps":
            torch.mps.empty_cache()
        elif device.type == "cuda":
            torch.cuda.empty_cache()

非原地更新的替代方案 #

如果需要保留函数返回新张量的逻辑，可显式删除旧张量帮助GC回收：

for i in range(10000):
    new_p1 = leak_example(p1, device)
    del p1  # 显式删除旧张量，加速内存回收
    p1 = new_p1
    if i % 100 == 0:
        gc.collect()
        if device.type == "mps":
            torch.mps.empty_cache()
        elif device.type == "cuda":
            torch.cuda.empty_cache()

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