切换 GPU 设备会影响 PyTorch 反向传播中的梯度吗？

Question

我用的是 Pytorch。在计算中，我在 GPU 中移动了一些数据和运算符 A。在中间步骤中，我将数据和运算符 B 移动到 CPU 并继续前进。

我的问题是：

我的算子B非常耗内存，无法在GPU中使用。这会影响（某些部分在 GPU 中计算，其他部分在 CPU 中计算）反向传播吗？

Answer 1

Pytorch 跟踪张量的位置。如果您使用 .cpu() 或 .to('cpu') pytorch 的本机命令，您应该没问题。

例如，参见 this model parallel 教程 - 计算在两个不同 GPU 设备之间进行拆分。

Answer 2

如果您的模型适合 GPU 内存，您可以让 PyTorch 在 DataParallel（单进程多线程）或 DistributedDataParallel（多进程多线程，单或多个节点）框架。

下面的代码检查您是否有 GPU 设备 torch.cuda.device_count() > 1 并设置 DataParallel 模式 model = nn.DataParallel(model)

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)

DataParallel 将相同的模型复制到所有 GPU，其中每个 GPU 消耗输入数据的不同分区，它可以显着加快训练过程，但它不适用于模型的某些用例太大而无法装入单个 GPU。

要解决此问题，您可以采用 model parallel 方法，将单个模型拆分到不同的 GPU 上，而不是在每个 GPU 上复制整个模型。

（例如，模型 m 包含 10 层：使用 DataParallel 时，每个 GPU 将拥有这 10 层中每一层的副本，而在使用时在两个 GPU 上并行建模，每个 GPU 可以承载 5 层）

一个示例，其中 .to('cuda:0') 指示图层应放置的位置。

import torch
import torch.nn as nn
import torch.optim as optim


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

backward() 然后自动考虑位置。

model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html

Answer 3

这个片段表明当计算通过不同的设备时梯度被保留。

def change_device():
    import torch.nn as nn
    a = torch.rand((4, 32))
    m1 = nn.Linear(32, 32)
    cpu = m1(a)
    gpu = cpu.to(0)
    m2 = nn.Linear(32, 32).to(0)
    out = m2(gpu)
    loss = out.sum()
    loss.backward()
    print(m1.weight.grad)
    # works like magic
    """
    tensor([[ 0.7746,  1.0342,  0.8706,  ...,  1.0993,  0.7975,  0.3915],
        [-0.5369, -0.7169, -0.6034,  ..., -0.7619, -0.5527, -0.2713],
        [ 0.3607,  0.4815,  0.4053,  ...,  0.5118,  0.3713,  0.1823],
        ...,
        [ 1.1200,  1.4955,  1.2588,  ...,  1.5895,  1.1531,  0.5660],
        [-0.1582, -0.2112, -0.1778,  ..., -0.2245, -0.1629, -0.0799],
        [-0.4531, -0.6050, -0.5092,  ..., -0.6430, -0.4665, -0.2290]])
    """

修改此代码段，当张量从 gpu 移动到 cpu 时，梯度也会保留。

切换 GPU 设备会影响 PyTorch 反向传播中的梯度吗？

Will switching GPU device affect the gradient in PyTorch back propagation?

pytorch

gpu

cpu

gradient-descent

backpropagation