自己记录用的:
torch.save(net.module.state_dict(), os.path.join(args.save, 'model.pth'))
改为
torch.save(net.state_dict(), os.path.join(args.save, 'model.pth'))
这篇博客记录了如何将使用`net.module.state_dict()`保存模型更改为使用`net.state_dict()`。作者指出,当模型在多GPU环境下训练且使用DataParallel或DistributedDataParallel时,需要保存`net.module.state_dict()`,但在单GPU或CPU环境下,直接保存`net.state_dict()`即可。这一更改简化了模型保存的代码,并适用于大多数不涉及分布式训练的情况。
自己记录用的:
torch.save(net.module.state_dict(), os.path.join(args.save, 'model.pth'))
改为
torch.save(net.state_dict(), os.path.join(args.save, 'model.pth'))
您可能感兴趣的与本文相关的镜像
PyTorch 2.7
PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

被折叠的 条评论
为什么被折叠?