PyTorch Models

晨光ABC2023-11-27 12:30

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))