保存整个模型:torch.save(模型名 , '文件名')
加载整个模型:model = torch.load(PATH)
只保存训练好的权重(参数):torch.save(模型名.state_dict() , 'save.pt')
加载模型的参数:model.load_state_dict(torch.load("save.pt")) #model.load_state_dict()函数把加载的权重复制到模型的权重中去,只保存权重的这种方式更好。
-
model.load_state_dict(torch.load(PATH))
-
model.eval()
保存模型的推理过程的时候,只需要保存模型训练好的参数,使用torch.save()保存state_dict,能够方便模型的加载。因此推荐使用这种方式进行模型保存。
记住一定要使用model.eval()来固定dropout和归一化层,否则每次推理会生成不同的结果。