PyTorch中torch.load()加载张量时存在内存泄漏问题

发现使用PyTorch的torch.load()方法加载已保存的张量时存在内存泄漏问题：即使显式删除加载的对象并调用垃圾回收，内存占用也未如预期下降。

最小复现代码 #

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
import gc

def create_and_save_image():
    image = torch.randn(30, 384, 384)
    ckpt = {"image": image}
    torch.save(ckpt, "test_image.pt")

@profile
def load_and_del():
    ckpt = torch.load("test_image.pt")
    del ckpt
    gc.collect()

if __name__ == '__main__':
    create_and_save_image()
    load_and_del()

复现步骤 #

• 执行命令：python3 -m memory_profiler mem_leak.py
• 观察加载和删除张量前后的内存占用情况

预期行为 #

调用torch.load()时内存占用上升，删除对象并调用gc.collect()后内存占用下降。

实际行为 #

调用torch.load()时内存占用上升，但删除对象并调用gc.collect()后内存占用未下降，表明存在内存泄漏。

系统信息 #

• PyTorch版本：2.3.0
• Python版本：3.8.5
• 操作系统：Linux
• CUDA版本：12.0
• GPU：NVIDIA GeForce RTX A6000（48GB）

补充信息 #

使用memory_profiler对内存占用进行分析，确认内存未按预期释放。

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