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

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

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