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

Question

火炬中的一些操作是就地执行的。例如，像 += 这样的速记运算符。是否可以就地执行其他操作，例如softmax？我目前正在研究语言处理。该模型在大量词汇表上生成一长串概率分布。这个最终输出张量负责大约 60% 的分配内存。这是一个巨大的问题，因为我需要计算一个 softmax 并且将所需的内存加倍。这是问题的一个例子。我对张量 t 不感兴趣，只对它的 softmax 感兴趣：import numpy as npimport torchimport torch.nn.functional as Ft = torch.tensor(np.zeros((30000,30000))).cuda()  #allocates 6.71 GB of GPUsoftmax = F.softmax(t, 1)  #out of memory errordel t  #too late, program crashed即使以下也不起作用：F.softmax(torch.tensor(np.zeros((30000,30000))).cuda(), 1)

Answer 1

到目前为止，PyTorch 是不可能的。您可以尝试推出自己的GPU 内核，但我看到前面有麻烦（如果不是墙的话），这可能是此操作最初不可用的原因。

除了归一化需要减少之外，Softmax 可以很容易地并行应用。减少是非常重要的，它们可以是就地或并行的（请记住，原子的使用相当于并发但不是并行的操作）。这意味着您的就地操作无论如何都必须在幕后分配，这违背了目的，或者运行得非常慢。请考虑我的回答有点推测性，我不是 GPGPU 的专家，但我的观点是，廉价、快速和正确地解决至少是一个困难的问题。

话虽如此，如果您只关心 softmax，在我看来，您似乎在进行推理。也许对于您的应用程序来说，将 logits 移动到 CPU 并在那里运行 softmax 是一个可行的选择，而您的 GPU 已经在处理下一批？