关于pytorch模型的保存(save)和加载(load)

目前常见的有两种方法：

一、保存和加载整个模型：torch.save() + torch.load()

示例：

torch.save(model, '/path/xxx.pth')
model = torch.load('/path/xxx.pth')

 

二、仅保存和加载参数：torch.save() + torch.load_state_dict()

示例：

torch.save(pre-trained_model.state_dict(), "/path/xxx.pth")  # 低版本pytorch用这个

torch.save(pre-trained_model.state_dict(), "/path/xxx.pth", _use_new_zipfile_serialization=False)  # 高版本（>=1.7.0）pytorch用这个

model = Model()
model.load_state_dict(torch.load('/path/xxx.pth'))
model.eval()  # 测试时添加