PyTorch内存管理与GPU-CPU迁移内存异常问题咨询

问题1：该行为是否为PyTorch的预期表现？ #

这是PyTorch的预期行为，根源在于PyTorch的CPU内存分配策略和Python垃圾回收的特性：

• 模型从GPU迁回CPU时，PyTorch会在CPU内存中开辟新的存储区域存放参数，这部分内存会被存入PyTorch的CPU内存池，而非直接还给操作系统——这种设计是为了避免频繁内存分配/释放带来的性能损耗。
• 前几次迭代内存持续增长，是因为内存池在逐步扩容至能满足需求的大小；3-4次后内存不再增长，说明内存池已稳定，后续会复用已分配的内存块，不会再向操作系统申请新内存。
• 首次加载内存低于后续迭代，是因为第一次运行时内存池尚未建立，后续迭代内存池已预留足够空间，看似占用更高，实则是内存池的预留空间，并非内存泄漏。

问题2：如何在不关闭线程的情况下释放CPU上的PyTorch内存？ #

可以通过以下几种方式处理：

1. 显式释放PyTorch CPU内存池（推荐PyTorch 1.10+） #

PyTorch提供了底层接口来清空CPU内存池中的闲置内存，无需关闭线程：

# 在需要释放内存的位置调用
from torch._C import _malloc_delete_all
_malloc_delete_all()

2. 优化内存使用流程 #

• 尽量避免反复在CPU和GPU之间迁移模型，如果必须迁移，确保迁移后立即清理旧的张量引用，不给垃圾回收留障碍。
• 保持模型相关逻辑在局部作用域内（你的代码已经做到了），避免全局变量持有模型或张量的引用，让Python垃圾回收能更高效地工作。

3. 结合激进的垃圾回收操作 #

在删除模型后，先清理所有参数引用，再触发垃圾回收和内存池释放：

@profile
def run_test():
    model = LargeNet()
    model = model.to('cuda')
    model = model.to('cpu')
    
    # 先清理参数引用
    for param in model.parameters():
        del param
    # 删除模型
    del model
    
    # 强制垃圾回收
    gc.collect()
    # 释放CPU内存池
    from torch._C import _malloc_delete_all
    _malloc_delete_all()
    torch.cuda.empty_cache()

注意事项 #

• 内存池的存在是为了提升性能，除非内存紧张到影响其他程序运行，否则不建议频繁强制释放——反复释放再分配会显著降低程序运行效率。
• 你观察到的内存不再增长是正常的稳定状态，后续运行不会再额外占用系统内存。

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