pytorch GPU内存管理

1、PyTorch 提供了 memory_allocated() 和 max_memory_allocated() 用于监视 tensors 占用的内存； memory_cached() 和 max_memory_cached() 用于监视缓存分配器所管理的内存.
2、PyTorch 提供了 empty_cache() l来释放所有未使用的缓存的内存，以便其它 GPU 应用能够使用. 但是，并不能释放 tensors 所占用的 GPU 显存，因此，其并不能增加 PyTorch 的可用 GPU 显存量.
pytorch 官网说明

• 查看特定设备上tensors当前所占用的 GPU 显存以字节bytes为单位

torch.cuda.memory_allocated(device=None)

• 查看特定设备上 tensors 最大所占用的 GPU 显存以字节bytes为单位，默认返回当前程序从开始所占用的最大显存.

torch.cuda.max_memory_allocated(device=None)

• 查看特定设备上缓存分配器当前所占用的 GPU 显存以字节bytes为单

torch.cuda.memory_cached(device=None)

• 查看特定设备上缓存分配器最大所占用的 GPU 显存以字节bytes为单位

torch.cuda.max_memory_cached(device=None)

• 释放缓存分配器当前所管理的所有未使用的缓存

torch.cuda.empty_cache()

例子：

try:
    output = model(input)
except RuntimeError as exception:
    if "out of memory" in str(exception):
        print("WARNING: out of memory")
        if hasattr(torch.cuda, 'empty_cache'):
            torch.cuda.empty_cache()
    else:
        raise exceptio

另外一点是，在测试时，避免忘记设置 torch.no_grad(). 如：

with torch.no_grad():
    inputs = None
    output = model(inputs)