[人工智能]pytorch搭建深度学习模型，哪些步骤占用显存

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