Pytorch训练时显存分配过程探究

2021-11-19 11:32:19 阅读：444 来源： 互联网

标签：显存 1024 torch 探究 Pytorch cuda memory print grad

参考：

https://blog.csdn.net/qq_37189298/article/details/110945128

========================================

代码：

import torch 
from torch import cuda 
import time
 
x = torch.zeros([1,1024,1024,128*2], requires_grad=True, device='cuda:0') 

print("1", cuda.memory_allocated()/1024**2)  

y = 5 * x 
# y.retain_grad()
print("2", cuda.memory_allocated()/1024**2)  


torch.mean(y).backward()     
print("3", cuda.memory_allocated()/1024**2)    
print(cuda.memory_summary())


time.sleep(60)

可以看到pytorch占显存共4777MB空间，其中变量及缓存共占4096空间。可以知道其中1024MB空间为缓存，可以手动释放，改代码：

import torch 
from torch import cuda 
import time
 
x = torch.zeros([1,1024,1024,128*2], requires_grad=True, device='cuda:0') 

print("1", cuda.memory_allocated()/1024**2)  

y = 5 * x 
# y.retain_grad()
print("2", cuda.memory_allocated()/1024**2)  


torch.mean(y).backward()     
print("3", cuda.memory_allocated()/1024**2)    


torch.cuda.empty_cache()
print(cuda.memory_summary())



time.sleep(60)

根据参考文章可知，1024*3MB是变量内存，其余700MB为其他内存，其中变量内存中有1024为x.grad，而且程序运行过程中显存分配峰值为4096MB，如下图：

其中包括 x.grad 和 y.grad 各1024MB空间。

如果保存非叶子节点的grad值，即保存y.grad，运行：

import torch 
from torch import cuda 
import time
 
x = torch.zeros([1,1024,1024,128*2], requires_grad=True, device='cuda:0') 

print("1", cuda.memory_allocated()/1024**2)  

y = 5 * x 
y.retain_grad()
print("2", cuda.memory_allocated()/1024**2)  


torch.mean(y).backward()     
print("3", cuda.memory_allocated()/1024**2)    


torch.cuda.empty_cache()
print(cuda.memory_summary())



time.sleep(60)

发现显存不够用了，也就是说保存y.grad后整体显存已经快达到5.9GB了，于是相同代码再Titan上运行：

发现总显存：

运行结果：

================================================

标签：显存,1024,torch,探究,Pytorch,cuda,memory,print,grad
来源： https://www.cnblogs.com/devilmaycry812839668/p/15576357.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

Pytorch训练时显存分配过程探究