背景 在公司用多卡训练模型,得到权值文件后保存,然后回到实验室,没有多卡的环境,用单卡训练,加载模型时出错,因为单卡机器上,没有使用DataParallel来加载模型,所以会出现加载错误。 原因 DataParallel包装的模型在保存时,权值参数前面会带有module字符,然而自己在单卡环境下,没有用DataParallel包装的模型权值参数不带module。本质上保存的权值文件是一个有序字典。 解决方法 1.在单卡环境下,用DataParallel包装模型。 2.自己重写Load函数,灵活。 from collections import OrderedDict def myOwnLoa
2022-03-19 14:03:37 46KB c OR pytorch
1
主要介绍了解决pytorch多GPU训练保存的模型,在单GPU环境下加载出错问题,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧
2021-12-28 19:55:06 45KB pytorch GPU训练保存 GPU环境 加载出错
1