Pytorch—模型的保存和加载torch.save()、torch.load()

1、模型保存

pytorch中**torch.save()**用于模型的存储。语法格式如下：

torch.save(obj,f,pickle_module=pickle,pickle_protocol=DEFAULT_PROTOCOL,_use_new_zipfile_serialization=True)

• obj：需要保存的对象，可以是整个模型或者是模型参数
• f：保存模型的路径
• pickle_module：用于清除元数据和对象的模块
• pickle_protocol：可以指定覆盖默认协议

模型保存有两种方式：

1、保存整个模型
同时保存整个神经网络的的模型结构和参数，保存的对象就是整个模型，此时保存的文件后缀名一般为.pt或.pth。

torch.save(model, ‘net.pkl’)

2、保存训练好的权重
只保存神经网络的训练模型的参数，保存的对象是model.state_dict()，此时保存的文件后缀名一般为.pkl或.pth。

torch.save(model.state_dict(), ‘net_params.pt’)

2、模型加载

pytorch中**torch.load()**用于模型的加载。语法格式如下：

torch.load(f,map_location=None,pickle_module=pickle,**pickle_load_args)

• f：导入模型文件
• map_location：指定如何重新映射存储位置的字符串或字典
• pickle_module：用于清除元数据和对象的模块
• pickle_load_args_：选择指定的参数传入

加载整个模型：

model = torch.load(‘net.pkl’) 

加载模型参数：

model.load_state_dict(torch.load(‘net_params.pt’))