火炬就地操作以节省内存(softmax)

火炬中的一些操作是就地执行的。例如,像 += 这样的速记运算符。


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


哈士奇WWW
浏览 236回答 2
2回答

繁花如伊

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

相关分类

Python