是否可以在 CPU 和 GPU 上存储一些张量以在 PyTorch 中训练神经网络？答案

【问题标题】：Is it possible to store some tensors on CPU and other on GPU for training neural network in PyTorch?是否可以在 CPU 和 GPU 上存储一些张量以在 PyTorch 中训练神经网络？
【发布时间】：2021-11-24 13:04:00
【问题描述】：

我在 PyTorch 中设计了一个神经网络，它需要大量的 GPU 内存或者以非常小的批量运行。

导致GPU Runtime错误由于三行代码，其中存储了两个新张量并进行了一些操作。

我不想以小批量运行我的代码。因此，我想在 CPU 上执行这三行代码（并因此存储这些新张量），并像往常一样将所有其他代码保留在 GPU 上。

可以吗？

【问题讨论】：

标签： pytorch gpu tensor

【解决方案1】：

这是可能的。
您可以使用命令.to(device=torch.device('cpu') 将相关张量从 GPU 移动到 CPU，然后再返回到 GPU：

orig_device = a.device  # store the device from which the tensor originated
# move tensors a and b to CPU
a = a.to(device=torch.device('cpu')) 
b = b.to(device=torch.device('cpu'))
# do some operation on a and b - it will be executed on CPU
res = torch.bmm(a, b)
# put the result back to GPU
res = res.to(device=orig_device)

几点说明：

在设备之间或 GPU 和 CPU 之间移动张量并不是一个不寻常的事件。用于描述它的术语是 "model parallel" - 您可以在 Google 上搜索更多详细信息和示例。
注意.to()操作是not an "in place" operation。
在 GPU 和 CPU 之间来回移动张量需要时间。在这种情况下，使用这种类型的“模型并行性”可能不值得。如果您正在为 GPU 空间而苦苦挣扎，您可以考虑使用gradient accumulation

【讨论】：