pytorch显存管理

PyTorch使用缓存分配器来加速内存分配。允许在不进行设备同步的情况下快速重新分配内存。缓存分配器中未被占用得内存，用nvidia-smi查看也显示为使用

torch.cuda.empty_cache()
释放缓存分配器中未被占用得显存（减少显存碎片），从而可被其他gpu服务使用，也可使用nvidia-smi命令查看

现象：

在部署gpu服务后，由于每次得请求数据大小不一致，随着请求数据变大（例如，文本长度变长，batchsize变大），gpu服务的缓存会越来越大。可使用torch.cuda.empty_cache()释放显存。但是，这样每次需要显存时，都要重新分配，比较耗时。

解决方法：

定时清缓存

if clear_cache_num == 5000:
    torch.cuda.empty_cache()
    clear_cache_num = 0

其他

torch.cuda.memory_allocated(device=None)
返回给定设备device的张量所占用的当前GPU内存

torch.cuda.max_memory_allocated(device=None)
返回给定设备device的张量所占用的GPU内存的最大值（从程序运行开始）

torch.cuda.memory_reserved(device=None)
返回给定设备device的缓存分配器管理的当前GPU内存

 torch.cuda.max_memory_reserved(device=None)
返回由缓存分配器管理的给定设备device的最大GPU内存（从程序运行开始）

cuda-memcheck工具可用于分析显存错误，为了避免缓存带来的影响，设置 PYTORCH_NO_CUDA_MEMORY_CACHING=1来关闭缓存分配器

cuda-memcheck工具:
https://docs.nvidia.com/cuda/cuda-memcheck/index.html

torch.cuda.empty_cache() 例子

import torch
import os
import time

device = torch.device('cuda:0')
# 定义两个tensor
tensor_4 = torch.randn(120, 3, 512, 512).float().to(device)
tensor_5 = torch.randn(80, 3, 512, 512).float().to(device)

time.sleep(1)
print("nvidia-smi start")
os.system('nvidia-smi -i 0')  # 1163M

gpu_memory_bf = torch.cuda.memory_allocated(device=device)
print("the gpu memory before tensor to cpu:{}".format(gpu_memory_bf))  # 629145600

# 然后释放，转到cpu
tensor_4 = tensor_4.cpu()
tensor_5 = tensor_5.cpu()
gpu_memory_af = torch.cuda.memory_allocated(device=device)
print("the gpu memory after tensor to cpu:{}".format(gpu_memory_af))  # 0
# 这里虽然将上面的显存释放了，但是我们通过Nvidia-smi命令看到显存依然在占用

time.sleep(1)
print("nvidia-smi before torch.cuda.empty_cache")
os.system('nvidia-smi -i 0')  # 1163M

torch.cuda.empty_cache()
# 只有执行完上面这句，显存才会在Nvidia-smi中释放

time.sleep(1)
print("nvidia-smi after torch.cuda.empty_cache")
os.system('nvidia-smi -i 0')  # 563M
# 好像执行完torch.cuda.empty_cache()，只回收一部分