pytorch的模型保存和读取
torch.save: Saves a serialized object to disk. This function uses Python’s pickle utility for serialization. Models, tensors, and dictionaries of all kinds of objects can be saved using this function.
torch.load: Uses pickle’s unpickling facilities to deserialize pickled object files to memory. This function also facilitates the device to load the data into (see Saving & Loading Model Across Devices).
torch.nn.Module.load_state_dict: Loads a model’s parameter dictionary using a deserialized state_dict. For more information on state_dict, see What is a state_dict?.
torch.save:将序列化的对象保存到磁盘。此函数使用Python的 pickle实用程序进行序列化。使用此功能可以保存各种对象的模型,张量和字典。
torch.load:使用pickle的解腌功能将腌制的目标文件反序列化到内存中。
torch.nn.Module.load_state_dict:使用反序列化的state_dict加载模型的参数字典 。
保存权重和加载权重
save
torch.save(model.state_dict(), PATH)
load
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
保存/加载整个模型
保存:
torch.save(model, PATH)
加载:
model = torch.load(PATH)
model.eval()
保存和加载超参数
保存:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)
加载:
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
迁移学习的加载模型
部分加载模型或加载部分模型
方法1:先加载参数,再改掉模型
model = TheModelBClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH), strict=False)
strict=False表示可以不准确的匹配,以忽略不匹配的键
inchannel = model.fc.in_features # 这里拿出来网络中fc层的输入层数
model.fc = nn.Linear(inchannel, 5) # 把网络里面的fc层改掉
方法2:
torch.load(PATH)#先把权重载入到内存中
在把要更改的层删掉
再在载入到模型修改好的模型中
这两种方式都行
节省CPU,直接将权重载入到GPU
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want
model.to(device)