pytorch中训练模型的保存和提取

有两种方式

只保存模型的参数：

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

、
之后需要的时候把模型参数再次提取出来：

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

保存整个模型

torch.save(the_model, PATH)

提取时：

the_model = torch.load(PATH)