torch.load(), 解决加载模型时出现NoModule named ‘models‘...

加载模型时,突然报出这种错误,大概率是因为使用torch.save()保存模型时,保存的是完整的模型,包括模型结构和权重参数。

要使用torch.load()加载完全保存的模型,可能会因为你加载模型的脚本所在路径与原始保存模型时脚本所在路径不一致,导致在加载模型时会报以上错误。

解决办法:

1.把加载模型的脚本放到当时保存模型时的路径下去执行加载。

2.

model = Model()
m_module = torch.jit.script(model)
torch.jit.save(m_module, 'total_model.pth')
torch.jit.load('total_model.pth')

3. 在保存模型时,只保存字典,不要保存整个模型了。

torch.save(model.state_dict(), PATH)

这种方式在加载模型时,需要先加载模型结构。