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

【问题标题】：Will switching GPU device affect the gradient in PyTorch back propagation?切换 GPU 设备会影响 PyTorch 反向传播中的梯度吗？
【发布时间】：2021-10-23 11:40:45
【问题描述】：

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

我的问题是：

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

【问题讨论】：

标签： pytorch gpu cpu gradient-descent backpropagation

【解决方案1】：

如果您的模型适合 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

【讨论】：

【解决方案2】：

Pytorch 会跟踪张量的位置。如果你使用 .cpu() 或 .to('cpu') pytorch 的原生命令应该没问题。

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

【讨论】：