Pytorch:当将张量移动到GPU时，内存会发生什么

Question

我试图了解当张量被发送到GPU时，RAM和GPU内存会发生什么。

在下面的代码示例中，我创建了两个张量-大张量arr=torch。张量ones((1000010000((和小张量c=torch。张量.ones(1(。张量c被发送到目标函数步骤内的GPU，该步骤由多处理调用。水塘在这样做的过程中，每个子进程在GPU上使用487 MB，RAM使用量达到5 GB。请注意，大张量arr在调用Pool之前只创建过一次，而不是作为参数传递给目标函数。当所有东西都在CPU上时，Ram的使用量不会激增。

关于这个例子，我有以下问题：

我在送火炬。Tensor.ones(1(到GPU，但它消耗了487MB的GPU内存。即使底层张量很小，CUDA是否也会在GPU上分配最小量的内存？GPU内存对我来说不是问题，这只是为了让我了解如何进行分配。

问题出在RAM的使用上。尽管我正在向GPU发送一个小张量，但似乎内存中的所有内容(大张量arr(都被复制到每个子进程(可能是固定内存(。那么，当张量被发送到GPU时，哪些对象被复制到固定内存？我在这里错过了一些东西，因为当我只发送一个特定的对象时，准备所有要发送到GPU的东西是没有意义的。

谢谢！

from multiprocessing import get_context
import time
import torch
dim = 10000
sleep_time = 2
npe = 4  # number of parallel executions
# cuda
if torch.cuda.is_available():
dev = 'cuda:0'
else:
dev = "cpu"
device = torch.device(dev)

def step(i):
c = torch.ones(1)
# comment the line below to see no memory increase
c = c.to(device)
time.sleep(sleep_time)

if __name__ == '__main__':
arr = torch.ones((dim, dim))
# create list of inputs to be executed in parallel
inp = list(range(npe))
# sleep added before and after launching multiprocessing to monitor the memory consumption
print('before pool')  # to check memory with top or htop
time.sleep(sleep_time)
context = get_context('spawn')
with context.Pool(npe) as pool:
print('after pool')  # to check memory with top or htop
time.sleep(sleep_time)
pool.map(step, inp)
time.sleep(sleep_time)

Answer 1

我正在向GPU发送torch.Tensor.ones(1)，但它消耗了487 MB的GPU内存。即使底层张量很小，CUDA是否也会在GPU上分配最小量的内存？

CUDA设备运行时为上下文建立时的各种事情保留内存，其中一些是固定大小的，另一些是可变的，可以由API调用控制(有关更多信息，请参阅此处(。显式或延迟地在设备上建立上下文的第一个API调用将产生GPU内存消耗的跳跃，这是完全正常的。在这种情况下，我认为第一个张量的创建触发了内存开销的分配。这是CUDA运行时的属性，而不是PyTorch或张量。