解决nn.DataParallel失效卡死,无输出
pytorch中的nn.DataParallel可以进行单机多卡训练。如果确实有多个可用的显卡,
torch.cuda.device_count() >1
但程序卡死在跳转到自定义模型的forward()这一步,且没有任何输出,则可能是没有将模型变成module。
即在
model = MyNetwork()#自定义的神经网络模型
model = nn.DataParallel(model)
之后,还需
model = model.module
此外优化器(optimizer)也需要使用module。
optimizer = nn.DataParallel(optimizer)
optimizer.module.step()
2452





