Pytorch显存分配机制是如何具体运作的？

本文共计464个文字，预计阅读时间需要2分钟。

对于显存不足的PyTorch研究者来说，掌握PyTorch的显存分配机制是非常重要的。以下通过实验来展示PyTorch显存分配的过程。

实验代码如下：

pythonimport torchfrom torch import cuda

x=tif cuda.is_available(): x=x.cuda()

　　对于显存不充足的炼丹研究者来说，弄清楚Pytorch显存的分配机制是很有必要的。下面直接通过实验来推出Pytorch显存的分配过程。

　　实验实验代码如下：

import torch from torch import cuda x = torch.zeros([3,1024,1024,256],requires_grad=True,device='cuda') print("1", cuda.memory_allocated()/1024**2) y = 5 * x print("2", cuda.memory_allocated()/1024**2) torch.mean(y).backward() print("3", cuda.memory_allocated()/1024**2) print(cuda.memory_summary())

输出如下：

　　代码首先分配3GB的显存创建变量x，然后计算y，再用y进行反向传播。