代码在单卡训练时没有问题,但是在多卡训练(DP或者DDP模式)时,会在loss.backward()处报错,但是报错并不详细,只停留在loss.backward(),再往后就是pytorch后端C++代码了。报错如下:
Runtime Error: Function BroadcastBackward returned an invalid gradient at index 770- got [0] but expected shape compatible with [0,8,3,3]
可见是反向传播时shape不一致。但是经过检查,我的loss并没有出错,单卡训练测试也没问题,这个就比较棘手了。
以下是我的排查步骤。
1. 打印计算图
先安装计算图可视化的包:
pip install torchviz
pip install graphviz
sudo apt-get install graphviz
然后在loss处插入如下代码:
from torchviz import make_dot
# 先计算出loss
loss = model_loss(output, gt)
# 创建计算图
dot = make_dot(loss, params=dict(model.named_parameters()))
# 保存图像
dot.render("model_graph", format="svg")
loss.backward()
我并没有从计算图中发现什么问题。
2. 寻找没有参与反向传播的参数
使用单卡调试,在loss.backward()之后、optimizer.step()之前加入如下代码:
for name, param in model.named_parameters():
if param.grad is None:
print(name)
打印出来的参数就是没有参与loss运算的部分,它们梯度为None。
在这里,我发现有一些模块没有参与loss运算,原因是我定义了一些用于实验的模块,但是后续forward当中没有使用。例如,我定义了一个modulelist,但是只有前面3个模块被使用了,后面2个模块定义了但是没有使用,却也作为参数传入了(因为整个modulelist都传入了)。
经过修改,bug修复!我认为出现“单卡可以运行但是多卡报错”的原因是分布式计算时需要保证计算图在梯度回传时都要有梯度。