一、定义一个容易识别的网络
在正式介绍模型的保存和加载之前，我们首先定义一个基本的网络Net，它只包含一个全连接层：
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.layer = nn.Linear(1, 1)
self.layer.weight = nn.Parameter(torch.FloatTensor([[10]]))
self.layer.bias = nn.Parameter(torch.FloatTensor([1]))
def forward(self, x):
y = self.layer(x)
return y
二、保存Net的参数值
保存模型参数之前，需要知道 Net的参数值存储在其state_dict（状态字典）属性中，我们查看一下net的state_dict包含哪些参数：
print(net.state_dict())
我们将会得到 net包含的所有参数名称与参数值：
包含一个weight和一个bias，对应的值分别是10和1，和我们之前定义的全连接层一致。我们需要保存的就是这个state_dict，保存的函数为”torch.save()”， 参数是我们需要保存的dict和存储路径：
torch.save(obj=net.state_dict(), f="models/net.pth")
这有可能会报错TypeError: state_dict() missing 1 required positional argument: ‘self’
报错原因是上面定义的class是一个类，不能直接obj=,正确的形式是obj=Net()
现在，同级目录models下将会出现net.pth文件，pth文件中的内容就是net的参数名称和值对应的state_dict，如下：
三、加载Net参数值并用于新的模型
最后一个步骤就是从pth文件中重新获取Net参数值，并把参数值装载到新定义的Model对象中。这里我们重新定义一个结构和Net类相同的类Model， 区别仅仅是Model参数初始值和Net不同，代码如下：
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.layer = nn.Linear(1, 1)
self.layer.weight = nn.Parameter(torch.FloatTensor([[0]]))
self.layer.bias = nn.Parameter(torch.FloatTensor([0]))
def forward(self, x):
out = self.layer(x)
return out
这里将Model的初始值权重w和偏差都设置为0，查看其state_dict：
model = Model()
print(model.state_dict())
得到的w和b值与预期相同，均为0，如下：
现在，我们将model对象的参数值设置为net.pth中的值，需要 使用”model.load_state_dict()”函数重置model的参数值为”torch.load(models/ net.pth)”中的参数值，如下：
model.load_state_dict(torch.load("models/net.pth"))
print(model.state_dict())
至此，model的w和b值就不再是0了，而是net中w和b对应的10和1，如下：
其中参数值重载的核心函数为”model.load_state_dict()”，每个继承自nn.Module的网络都能通过这个函数设定参数值。
四、优化器与epoch的保存
保存优化器参数值和epoch值的主要目的是用于继续训练，保存的流程依旧是先”torch.save()”再”torch.load_state_dict()”，我们首先定义一个Adam优化器、一个任意的epoch值与net如下：
net = Net()
Adam = optim.Adam(params=net.parameters(), lr=0.001, betas=(0.5, 0.999))
epoch = 96
现在， 创建一个字典来保存所有的对象，并用save函数保存这个字典：
all_states = {"net": net.state_dict(), "Adam": Adam.state_dict(), "epoch": epoch}
torch.save(obj=all_states, f="models/all_states.pth")
所有的对象都被保存到models文件夹下了：
可以使用load()函数把所有的对象再次提取出来：
reload_states = torch.load("models/all_states.pth")
print(reload_states)
五、总结
pytorch中state_dict()和load_state_dict()函数配合使用可以实现状态的获取与重载，load()和save()函数配合使用可以实现参数的存储与读取。其中最重要的部分是”字典”的概念，因为参数对象的存储是需要”名称”——”值”对应（即键值对），读取时也是通过键值对读取的。
Original: https://blog.csdn.net/m0_67635646/article/details/124349777
Author: m0_67635646
Title: pytoch中class定义神经网络的参数保存与加载
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/690489/
转载文章受原作者版权保护。转载请注明原作者出处！