pytorch_trick(4) 模型本地保存与读取方法

模型本地保存与读取方法

同时,借助state_dict()方法,我们可以实现模型或优化器的本地保存于读取。此处以模型为例,优化器的本地保存相关操作类似。

对于模型而言,其实也有state_dict()方法。通过该方法的调用,可以查看模型全部参数信息。

值得注意的是,模型的训练和保存,本质上都是针对模型的参数。而模型的state_dict()则包含了模型当前全部的参数信息。因此,保存了模型的state_dict()就相当于是保存了模型。

python 复制代码
# 设置随机数种子
torch.manual_seed(24)  

# 实例化模型  
tanh_model1 = net_class2(act_fun= torch.tanh, in_features=5, BN_model='pre')
tanh_model1.state_dict()

1、保存模型参数

首先,我们可以将该存有模型全部参数信息的字典对象赋给某个变量。

python 复制代码
t1 = tanh_model1.state_dict()
t1

其次,我们也可以通过torch.save来将该参数保存至本地。

python 复制代码
torch.save(tanh_model1.state_dict(), 'tanh1.pt')

对于torch.save函数来说,第一个参数是需要保存的模型参数,而第二个参数则是保存到本地的文件名。一般来说可以令其后缀为.pt.pth。而当我们需要读取保存的参数结果时,则可以直接使用load_state_dict方法。该方法的使用我们稍后就会谈到。

接下来进行模型训练,也就是模型参数调整。回顾此前学习内容,当我们进行模型训练时,实际上就是借助损失函数和反向传播机制进行梯度求解,然后利用优化器根据梯度值去更新各线性层参数。

python 复制代码
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(tanh_model1.parameters(), lr=0.05)
for X, y in train_loader:
    yhat = tanh_model1.forward(X)
    loss = criterion(yhat, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

训练完一轮之后,我们可以查看模型状态:

python 复制代码
tanh_model1.state_dict()

我们发现模型的参数已经发生了变化。当然,此时t1也随之发生了变化

python 复制代码
t1

2、还原模型参数

此时,如果我们想还原tanh_model1中原始参数,我们只能考虑通过使用load_state_dict方法,将本次保存的原模型参数替换当前的tanh_model1中参数,具体方法如下:

python 复制代码
torch.load('tanh1.pt')
tanh_model1.load_state_dict(torch.load('tanh1.pt'))
tanh_model1.state_dict()

至此,我们就完成了模型训练与保存的基本过程。当然,除了模型可以按照上述方法保存外,优化器也可以类似进行本地存储。

当然,结合此前介绍的深拷贝的相关概念,此处我们能否通过深拷贝的方式将模型参数保存在当前操作空间内然后再替换训练后的模型参数呢?同学们可以自行尝试

相关推荐
点云SLAM3 小时前
Eigen 中矩阵的拼接(Concatenation)与 分块(Block Access)操作使用详解和示例演示
人工智能·线性代数·算法·矩阵·eigen数学工具库·矩阵分块操作·矩阵拼接操作
木枷4 小时前
NAS-Bench-101: Towards Reproducible Neural Architecture Search
人工智能·物联网
BAOYUCompany4 小时前
暴雨服务器更懂人工智能+
运维·服务器·人工智能
飞哥数智坊4 小时前
Coze实战第17讲:工资条自动拆分+一对一邮件发送
人工智能·coze
cwn_4 小时前
自然语言处理NLP (1)
人工智能·深度学习·机器学习·自然语言处理
点云SLAM5 小时前
PyTorch中flatten()函数详解以及与view()和 reshape()的对比和实战代码示例
人工智能·pytorch·python·计算机视觉·3d深度学习·张量flatten操作·张量数据结构
智海观潮5 小时前
Unity Catalog与Apache Iceberg如何重塑Data+AI时代的企业数据架构
大数据·人工智能·ai·iceberg·catalog
爱分享的飘哥5 小时前
第三篇:VAE架构详解与PyTorch实现:从零构建AI的“视觉压缩引擎”
人工智能·pytorch·python·aigc·教程·生成模型·代码实战
柏峰电子5 小时前
市政道路积水监测系统:守护城市雨天出行安全的 “智慧防线”
大数据·人工智能·安全
蓑雨春归5 小时前
自主智能Agent如何重塑工作流自动化:技术、经济与未来展望
人工智能·chatgpt·自动化