Overview
pth模型保存时是按照"整个模型保存"和"只保存模型参数"会影响模型的加载和访问方式
python
torch.save(vgg16, "vgg16.pt")
torch.save(vgg16,"vgg16.ckpt")
torch.save(vgg16,"vgg16.pth")
torch.save(vgg16,"vgg16.pkl")Save & Load Models
保存整个模型
python
torch.save(model, PATH)
python
import torch
if __name__ == '__main__':
model_pth = r'D:\${modelPath}\${modelName}.pth'
net = torch.load(model_pth, map_location=torch.device('cpu'))
for key, value in net["state_dict"].items():
print(key,value.size(),sep=" ")只保存模型参数
python
torch.save(net.state_dict(),path2)
python
model.load_state_dict(torch.load(path2))