torch保存、读取模型

方法一:

将整个模型保存、读取:

torch.save(model, PATH)
model = torch.load(PATH)

方法二:

仅保存模型的state_dict()、读取:

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

这时仅读取了state_dict，还需要:

new_model.load_state_dict(torch.load(PATH)) # 其中的new_model是新初始化的模型

问题

靠，直接存模型翻车了，爆了类似于这样的错误，以后还是只保存state_dict吧。