1. 首页
  2. 编程语言
  3. Python
  4. 弄清Pytorch显存的分配机制

弄清Pytorch显存的分配机制

上传者: 2022-01-31 11:56:43上传 PDF文件 117.71 KB 热度 16次

  对于显存不充足的炼丹研究者来说,弄清楚Pytorch显存的分配机制是很有必要的。下面直接通过实验来推出Pytorch显存的分配过程。可以看到,创建x后与计算y后分别占显存3GB与6GB,这是合理的。另外,后面通过backward(),计算出x.grad,占存与x一致,所以最终一共占有显存9GB,这也是合理的。但是,输出显示了显存的峰值为12GB,这多出的3GB是怎么来的呢?首先画出计算图:下面通过列表的形式来模拟Pytorch在运算时分配显存的过程:  如上所示,由于需要保存反向传播以前所有前向传播的中间变量,所以有了12GB的峰值占存。  我们可以不存储计算图中的非叶子结点,达到节省显存的目的,即可以把上面的代码中的y=5*x与mean写成一步: 占显存量减少了3GB:

下载地址
用户评论