torch 中的一些操作是就地执行的。速记运算符,例如 +=。
其他操作是否可以就地执行,例如softmax
?
我目前正在从事语言处理工作。该模型在大量词汇上生成一长串概率分布。最终输出张量约占分配内存的 60%。这是一个很大的问题,因为我需要计算它的 softmax,这使得所需的内存加倍。
这是问题的一个示例。我对张量 t 不感兴趣,只对它的 softmax 感兴趣:
import numpy as np
import torch
import torch.nn.functional as F
t = torch.tensor(np.zeros((30000,30000))).cuda() #allocates 6.71 GB of GPU
softmax = F.softmax(t, 1) #out of memory error
del t #too late, program crashed
即使下面的方法也不起作用:
F.softmax(torch.tensor(np.zeros((30000,30000))).cuda(), 1)
最佳答案
我创建了一个就地版本的 softmax:
import numpy as np
import torch
import torch.nn.functional as F
# in-place version
t = torch.tensor(np.ones((100,200)))
torch.exp(t, out=t)
summed = torch.sum(t, dim=1, keepdim=True)
t /= summed
# original version
t2 = torch.tensor(np.ones((100,200)))
softmax = F.softmax(t2, 1)
assert torch.allclose(t, softmax)
回答我的问题:如果您想要就地函数,您必须通过将低级操作插入在一起来自己创建它们:
- 许多功能,例如
torch.exp
可以指定一个可选的out
参数。 - 分配
t[idx] = some
就位 - 简写运算符
/=
、*=
、+=
、-=
已到位<
这需要仔细调试并且可能不直观:
t = t / summed #not in-place
t /= summed #in-place
我读到就地操作可能会产生渐变问题。我将使用此代码进行更多测试。
关于python - torch 就地操作以节省内存(softmax),我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/53732209/