Question

割炬中的某些操作是就地执行的。速记运算符，例如+ =。

是否可以就地执行其他操作，例如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)

Answer 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] = something就位
速记运算符/=，*=，+=，-=就位

这需要仔细调试，并且可能不直观：

t = t / summed  #not in-place
t /= summed  #in-place

我读过，就地操作会产生梯度问题。我将使用此代码进行更多测试。

Answer 2

截至目前，PyTorch尚无法使用。您可以尝试自己动手GPU kernel，但我看到前方有麻烦（如果不是一堵墙），这很可能是该操作最初无法使用的原因。

除标准化外，Softmax可以轻松并行应用，这需要简化。归约是不平凡的，它们可以就地或并行（请记住，原子的使用等于并发而不是非并行操作）。这意味着就地操作无论如何都必须在幕后进行分配，这会破坏目的，或者运行速度非常慢。请认为我的回答有些投机，我不是GPGPU专家，但我的观点是，至少要廉价，快速和正确地解决它是一个难题。

话虽如此，如果您只关心softmax，对我来说听起来就像您在做推断。也许对于您的应用程序来说，将logit移至CPU并在那里运行softmax是一个可行的选择，而您的GPU已经处理了下一批？

火炬就地操作以节省内存（softmax）

2 个答案: