> 技术文档 > 模型的存储、加载和部署_模型训练存储卷

模型的存储、加载和部署_模型训练存储卷


定义损失函数并以此训练和评估模型

存储模型可以只存储state_dict或模型参数,每当需要部署经过训练的模型时,创建模型的对象并从文件加载参数,这是 Pytorch 创建者推荐的方法。

目录

模型的存储、加载

模型的部署


模型的存储、加载

承接上文,完成模型的训练后,需要将训练的参数存储在文件中,以供部署和使用。

#定义路径path2weights=\"./models/weights.pt\"#将state_dict存储到文件torch.save(model.state_dict(), path2weights)

为了从文件中加载模型参数,定义一个 Net 的对象类并加载state_dict

#定义随机初始权重模型_model = Net()#加载文件中的state_dictweights=torch.load(path2weights)#赋予权重_model.load_state_dict(weights)

加载成功如下 

模型的部署

将模型加载到内存中后,可以将新数据传递给模型

import matplotlib.pyplot as plt#抽取一个n=10张量n=10x= x_val[n]y=y_val[n]print(x.shape)plt.imshow(x.numpy()[0],cmap=\"gray\")

对张量进行预处理

#将维度扩展为 1*C*H*Wx= x.unsqueeze(0)#转换为torch.float32格式x=x.type(torch.float)

得到模型预测 

#获取模型输出output=_model(x)#获取预测结果pred = output.argmax(dim=1, keepdim=True)print (pred.item(),y.item())

 

易装网