解决PyTorch 出现RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 问题

解决PyTorch 出现RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 问题

在使用pytorch框架训练中，进行loss.backward()一步时可能出现如下的问题：

RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type TensorOptions(dtype=float, device=cuda:0, layout=Strided, requires_grad=false) but got TensorOptions(dtype=float, device=cpu, layout=Strided, requires_grad=false) (validate_outputs at /pytorch/torch/csrc/autograd/engine.cpp:484)
frame #0: c10::Error::Error(c10::SourceLocation, std::string const&) + 0x46 (0x7f4228517536 in /usr/local/lib/python3.8/dist-packages/torch/lib/libc10.so)
frame #1: + 0x2d84224 (0x7f418672b224 in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so)
frame #2: torch::autograd::Engine::evaluate_function(std::shared_ptrtorch::autograd::GraphTask&, torch::autograd::Node*, torch::autograd::InputBuffer&) + 0x548 (0x7f418672cd58 in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so)
frame #3: torch::autograd::Engine::thread_main(std::shared_ptrtorch::autograd::GraphTask const&, bool) + 0x3d2 (0x7f418672ece2 in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so)
frame #4: torch::autograd::Engine::thread_init(int) + 0x39 (0x7f4186727359 in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_cpu.so)
frame #5: torch::autograd::python::PythonEngine::thread_init(int) + 0x38 (0x7f422f0ad3d8 in /usr/local/lib/python3.8/dist-packages/torch/lib/libtorch_python.so)
frame #6: + 0xd6cb4 (0x7f422ff5acb4 in /usr/lib/x86_64-linux-gnu/libstdc++.so.6)
frame #7: + 0x9609 (0x7f423311e609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #8: clone + 0x43 (0x7f423325a103 in /lib/x86_64-linux-gnu/libc.so.6)

这个问题实际上是由于反向传播时的device不一致问题。训练过程中，前向传播时我们需要模型、输入位于同一设备上（都在cuda或都在cpu），如果这时出了问题能够清晰的找到报错信息，一个类型是Tensor，一个是CUDATensor。但是在反向传播时，这个问题不是很容易找到，反向传播要求loss与模型也位于同于设备上。即使是多卡并行，最终也会将Loss集中在一张卡上，这也是为什么pytorch原生nn.DataParallel 经常会导致所谓的多卡负载不平衡，有一张主卡要处理所有的loss的。

那问题来了，为什么loss会跟模型不在同一设备呢？按理说我们把input、model和label都通过.to(device)操作部署在GPU上，那model(input)得到的输出output也是在GPU上的，跟同在GPU的张量label计算loss，结果依然在GPU上，怎么就出问题了呢？

其实对于一般的loss确实不会有问题，但是对于多损失函数模型，比如loss = loss_a + loss_b +...，尤其是图像领域多尺度的损失，有可能通过一个for循环求解总的损失函数，自己新建了一个值为0.0的初始Tensor，通过for循环，每次叠加一个尺度下的损失，这时总的loss张量所在的device取决于新建的初始张量，如果忘了此时通过.to(device)部署在GPU上，那么即使每个尺度的损失都是通过两个CudaTensor计算出来的，叠加的总loss依然是cpu上的Tensor。此时，调用backward就会出现上述的错误。
例如下面这个损失，输入是元组input，其中的每个元素是不同scale下的模型输出，我们需要求不同尺度下输出与target的mse损失的加权和，如果这里的初始loss不加这个.cuda()那么最终返回的loss是一个cpu上的Tensor，与模型及输入不符，那么就会出现上述的错误信息。

  loss = torch.Tensor([0.0]).float().cuda()
  scale = 1
  weight = 1 / 2
  for i in range(0, 6):
    gt = F.interpolate(target, scale_factor=scale)
    loss += weight * F.mse_loss(input[i], gt)
    scale /= 2
    weight /= 2

如果不采用上面的.cuda()方法，我们也可以：

  scale = 1
  weight = 1 / 2
  gt = F.interpolate(target, scale_factor=scale)
  loss = weight * F.mse_loss(input[i], gt)
  for i in range(1, 6):
    scale /= 2
    weight /= 2
    gt = F.interpolate(target, scale_factor=scale)
    loss += weight * F.mse_loss(input[i], gt)

先计算scale=1（原始尺寸）下的loss，此时input、target全是cudaTensor，loss自然也是在GPU上，接下来叠加其他尺度的loss，也可以让最终的loss是cudaTensor而避免出错。