模型保存的两种类型torch.save

torch.save(model, 'model_face.pt')

一个常见的PyTorch函数是使用.pt文件扩展名来保存张量,

model是我训练后的模型

后面的参数'model_face.pt'就是我模型保存的类名,存放同一文件夹下,当然可以指定文件位置。

torch.save(model.state_dict(), './cnn_model.pt')

保存同样为pt格式,但存储的对应的是python的字典类型,这个是为了更好调用。

这两种训练后的模型model保存的方式比较常见,但需要注意的是,我们在调用模型来预测或者测试时,必须按照你原先保存的形式来调用,不然加载会出错的欧。

第一种(save(model, 'model_face.pt'))保存对应的加载方式:

torch.load("./model_face.pt")

第二种(save(model.state_dict(), './cnn_model.pt'))的加载方式:

model.load_state_dict(torch.load("./cnn_minist.pt"))

特别需要注意第二种的加载方式,很容易忘记里面的torch.load,呵呵哈哈哈

你可能感兴趣的:(pytorch学习笔记,python,神经网络,机器学习,深度学习)