关于PyTorch中模型保存与FineTuning的问题

这几天被一个问题卡了2天

在神经网络中加入了多头自注意力机制，训练有一定的效果，把模型保存下来之后，换了个loss继续训练做finetuning，然后就出现了这样的问题：

已经训练好的模型的loss在1左右，然后做Finetuing的时候，按理说加载了预训练模型，第一个Epoch的loss应该也在1左右（对同一个数据集），结果训练的时候第一个Epoch的loss飙升，升到了30多，这让我十分困惑，因为以前从来没有遇到过。

经过搜索资料我得知，由于我用的是Adam优化器，它基于动量进行优化，在训练的过程中也有自己的参数，在FineTuning的时候也要把他保存下来。代码如下：

    if avgloss < min_loss:
        print("reload best.pth")
        min_epoch = i
        min_loss = avgloss
        torch.save(SncNet.state_dict(), "best.pth")
        torch.save(optimizer.state_dict(),"Adam_best.pth")

在加载的时候也要把optimizer加载进去，代码如下：

learning_rate = 0.00000001
optimizer = torch.optim.Adam(SncNet.parameters(), lr=learning_rate)
if opt_pretrained:
    print("loading pretrained optimizer")
    pretrained_opt = torch.load("Adam_pretrained.pth")
    optimizer.load_state_dict(pretrained_opt)

解决之后，继续训练，然后又发现问题了——

我把每个iteration的loss打印下来，发现一个iteration的loss还是很高，但是第二个和第三个就开始骤降，降到了1左右（和预训练模型一样的loss水平）。所以问题还是没解决。

想到这个问题是加了多头自注意力机制的出现的，因此我把保存下来的模型打印出来，输出中并没有找到我所定义的msa，于是可以得出，MSA的参数没有保存。

sonicNet = NestedUNet_SA(num_classes=3).to("cuda")
sonicNet.load_state_dict(torch.load(modelpath))
for name,param in sonicNet.named_parameters():
    # if "msa" in name:
    # print("name:",name,"  param:",param)
    print("name:",name)

最后询问了ChatGPT，他给出的答案是：

问：pytorch中保存模型的时候，为什么没有把多头自注意力机制层保存下来

答：在 Pytorch 中，使用 torch.save() 方法来保存模型的 state_dict，其中 state_dict 是由层名称及其对应的参数组成的字典。这意味着只有模型中的权重和偏置（parameters）被保存下来，而模型中的其他内容，如多头自注意力机制层或其他不包含可学习参数的层（如 relu）并不会被保存。

这是因为在 PyTorch 中，只有包含可学习参数的层才被认为是模型的重要组成部分，需要被保存下来，因为它们是我们最终想要训练和优化的部分。另一方面，其他层（如 dropout、batchnorm 和激活函数）则不包含可学习参数，只是一些转换函数，这些函数会对输入数据进行变换，对于它们来说，它们的状态通常不需要被保存。

因此，如果您想要保存一个包含多头自注意力机制层的模型，您可以尝试保存整个模型，而不仅仅是它的 state_dict。例如，可以使用以下代码将整个模型保存到磁盘上：

# 定义并训练一个带有多头自注意力机制层的模型
model = MyModel()
...
# 保存整个模型
torch.save(model, 'my_model.pt')

当您想要加载整个模型时，只需要使用 torch.load() 方法来加载整个模型即可：

# 加载整个模型
model = torch.load('my_model.pt')