解决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而避免出错。