使用多gpu训练时
model = torch.nn.DataParallel(model, device_ids=[1, 2, 3, 4])
若模型采用多GPU训练,则在模型保存时:
torch.save(model.module.state_dict(), model_out_path)
若单GPU则:
torch.save(mode.state_dict(), model_out_path)
使用多gpu训练时
model = torch.nn.DataParallel(model, device_ids=[1, 2, 3, 4])
若模型采用多GPU训练,则在模型保存时:
torch.save(model.module.state_dict(), model_out_path)
若单GPU则:
torch.save(mode.state_dict(), model_out_path)