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

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

到目前为止，PyTorch 是不可能的。您可以尝试推出自己的GPU 内核，但我看到前面有麻烦（如果不是墙的话），这可能是此操作最初不可用的原因。除了归一化需要减少之外，Softmax 可以很容易地并行应用。减少是非常重要的，它们可以是就地或并行的（请记住，原子的使用相当于并发但不是并行的操作）。这意味着您的就地操作无论如何都必须在幕后分配，这违背了目的，或者运行得非常慢。请考虑我的回答有点推测性，我不是 GPGPU 的专家，但我的观点是，廉价、快速和正确地解决至少是一个困难的问题。话虽如此，如果您只关心 softmax，在我看来，您似乎在进行推理。也许对于您的应用程序来说，将 logits 移动到 CPU 并在那里运行 softmax 是一个可行的选择，而您的 GPU 已经在处理下一批？