当保存和加载模型时,需要熟悉三个核心功能:
torch.save:将序列化对象保存到磁盘。此函数使用Python的pickle模块进行序列化。使用此函数可以保存如模型、tensor、字典等各种对象。
torch.load:使用pickle的unpickling功能将pickle对象文件反序列化到内存。此功能还可以有助于设备加载数据。
torch.nn.Module.load_state_dict:使用反序列化函数 state_dict 来加载模型的参数字典。
Python中对于模型数据的保存和加载操作都是引用Python内置的pickle包,使用pickle.dump()和pickle.load()方法。在Pytorch中也有同样功能的方法提供。
>>>torch.save(model,'model.pkl') #保存整个模型
>>>model = torch.load('model.pkl') #加载整个模型
>>>torch.save(alexnet.