pytorch的显存机制,一直不明白,今天看了两篇文章,终于有些明白了,这两篇文章如下:
https://zhuanlan.zhihu.com/p/424512257
https://blog.csdn.net/qq_43827595/article/details/115722953
根据第一篇知乎文章及其评论,我自己写了一套代码来记录各个步骤中显存的占用情况,以下是通过jupyter notebook中转化而来得到的markdown文件内容:
import torch
import torch.nn as nn
class Net(nn.Module):
def __init__(self):
super().__init__()
self.linear1 = nn.Linear(1024,1024, bias=False)
self.linear2 = nn.Linear(1024, 1, bias=False)
def forward(self, x):
x = self.linear1(x)
x = self.linear2(x)
out = sum(x)
return out
# 输入数据
inputs = torch.tensor([[1.0]*1024]*1024).cuda() # +4194304
print(torch.cuda.memory_allocated())
4194304
# 模型初始化
net = Net().cuda() # +4194304+4096 前者是第一个线性层的参数所占显存,后者是第二个线性层所占显存
print(torch.cuda.memory_allocated())
8392704
# 正向传播
out = net(inputs) # +4194304+512,4194304是模型第一层输出所占显存之和,512是out所占显存
# 这里之所以不需要第二层的输出,是因为求导的时候用不到,pytorch在维护计算图的时候将其释放
# 每一层的输出,是否被舍弃,关键要看输出的结果是否在求导的表达式中
print(torch.cuda.memory_allocated())
12587520
# 后向传播
out.backward() # +4096,这个4096,是模型第二层中每个参数的梯度,
# 因为第二层的输出被舍弃掉了,需要新开辟显存来存储梯度
print(torch.cuda.memory_allocated())
12591616
from torch.optim import Adam
optimizer = Adam(net.parameters(), lr=1e-3) # +4198912 优化器参数所占显存
print(torch.cuda.max_memory_allocated())
16790528
optimizer.step() # +16785408 优化器相关参数所占显存
print(torch.cuda.max_memory_allocated())
29377024
16790528 - 12591616
4198912