1. 保存模型参数
torch.save(model.state_dict(), PATH)
使用多卡训练单卡加载时使用
torch.save(model.module.state_dict(), PATH)
2. 加载参数
model = mymodel() # 初始化模型
model.load_state_dict(torch.load(PATH))
加载时遇到报错
RuntimeError: CUDA out of memory. Tried to allocate …
解决方法:将权重load进入cpu,再让模型加载
model .load_state_dict(torch.load(PATH , map_location='cpu'))
model .to(device)