Pytorch模型加载和保存

一、模型的保存和加载

在模型的训练过程中，可能会因为各种原因停止，因此需要注意保存每一轮的epoch的模型（一般保存当前模型和最好模型）。

1.Pytorch保存和加载模型

1.1 官方推荐的方法

#第一种：只存储模型中的参数，该方法速度快，占用空间少（官方推荐使用）
model = VGGNet()
torch.save(model.state_dict(), PATH)          #存储model中的参数

new_model = VGGNet()                          #建立新模型
new_model.load_state_dict(torch.load(PATH))   #将model中的参数加载到new_model中

#第二种：存储整个模型
model = VGGNet()
torch.save(model, PATH)                       #存储整个模型

new_model = torch.load(PATH)                  #将整个model加载到new_model中
#new_model 不再需要第一种方法中的建立新模型的步骤

'''
关于上面表达式中PATH参数的说明：
PATH参数是你保存文件的路径，并且需要指定保存文件的文件名，如：
torch.save(model, '/home/user/save_model/checkpoint.pth')
即将该模型保存在/home/user/save_model路径下的checkpoint.pth文件中，保存的文件格式约定为.pth或.pt
new_model = torch.load('/home/user/save_model/checkpoint.pth')

但是在pytorch1.6版本中，torch.save存储的文件格式采用了新的基于压缩文件的格式 .pth.tar
torch.load依然保留了加载了旧格式.pth的能力
'''

1.2 保存checkpoint（检查点）

保存的训练模型信息不仅包含模型的参数信息，还可以包含其他信息，如当前的迭代次数，优化器的参数等，以便用于后面恢复训练。

#保存模型的状态，可以设置一些参数，后续可以使用
state = {'epoch': epoch + 1,#保存的当前轮数
         'state_dict': mymodel.state_dict(),#训练好的参数
         'optimizer': optimizer.state_dict(),#优化器参数,为了后续的resume
         'best_pred': best_pred#当前最好的精度
          ,....,...}

#保存模型到checkpoint.pth.tar
torch.save(state, ‘checkpoint.pth.tar’)
#如果是best,则复制过去
if is_best:
    shutil.copyfile(filename, directory + 'model_best.pth.tar')

checkpoint = torch.load('model_best.pth.tar')
model.load_state_dict(checkpoint['state_dict'])#模型参数
optimizer.load_state_dict(checkpoint['optimizer'])#优化参数
epoch = checkpoint['epoch']#epoch，可以用于更新学习率等

#有了以上的东西，就可以继续重新训练了，也就不需要担心停止程序重新训练

在断点处继续训练

这个方法可以用来提高模型最终的准确率。

训练的时候经常遇到一种情况就是，可能设置了500个epoch，但是当模型跑到200个epoch之后，损失和精度就上不去了，或者出现明显的过拟合了，剩下的300个epoch，纯属浪费时间，这300个epoch对模型几乎没啥影响。但是其实在200个epoch时保存参数和学习率，然后重新加载模型，并将学习率调低，损失有可能会再降一些，精度有可能会得到提高。

原文链接：pytorch模型保存与加载总结_pytorch加载模型-CSDN博客

1.3 不同设备上的模型存储与加载

#1、在CPU上存储模型，在GPU上加载模型
#CPU存储
torch.save(model.state_dict(), PATH)
#GPU加载
device = torch.device('cuda')
model = Model()
model.load_state_dict(torch.load(PATH, map_location='cuda:0')) #可以选择任意GPU设备
model.to(device)

----------------------------------------------------------------------------------
#2、在GPU上存储，CPU上加载
#GPU存储
torch.save(model.state_dict(), PATH)
#CPU加载
device = torch.device('cpu')
model = Model()
model.load_state_dict(torch.load(PATH, map_location=device))

----------------------------------------------------------------------------------
#3、在GPU上存储，在GPU上加载
#GPU存储
torch.save(model.state_dict(), PATH)
#GPU加载
device = torch.device('cuda')
model = Model()
model.load_state_dict(torch.load(PATH))
model.to(device)

----------------------------------------------------------------------------------
#4、存储和加载使用过torch.nn.DataParallel的模型

#（1）多卡训练，单卡加载部署

'''
这种情况要防止参数保存的时候没有加module,那么保存的参数名称是module.conv1.weight，
而单卡的参数名称是conv1.weight

，这时就会报错，找不到相应的字典的错误。
此时可以通过手动的方式删减掉模型中前几位的名称，然后重新加载。
不懂代码可以先看一下第2部分内容模型参数存储内容解析
'''

model = torch.nn.DataParallel(model)
#存储
torch.save(model.module.state_dict(), PATH)

#加载
kwargs={'map_location':lambda storage, loc: storage.cuda(gpu_id)}
def load_GPUS(model,model_path,kwargs):
    state_dict = torch.load(PATH, **kwargs)
    # create new OrderedDict that does not contain 'module.'
    from collections import OrderedDict
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[7:] # remove 'module.'
        new_state_dict[name] = v
    # load params
    model.load_state_dict(new_state_dict)
    return model

#（2）单卡训练，多卡加载部署

'''
此时唯有记住一点，因为单卡训练参数是没有module的，而多卡加载的参数是有module的，
因此需要保证参数加载在模型分发之前。
'''

#存储
torch.save(model.state_dict(), PATH)
#加载
model.load_state_dict(torch.load(PATH))
model = torch.nn.DataParallel(model)     #模型分发

#（3）多卡训练，多卡加载部署

'''
环境如果没有变化，则可以直接加载，如果环境有变化，则可以拆解成第1种情况，然后再分发模型。
'''

当使用多块GPU进行训练的时候，保存的参数名称是module.conv1.weight，而单卡的参数名称是conv1.weight，这时你保存完模型后，再重新加载就会报错，找不到相应的字典。

解决办法一

保存的时候把参数名中的‘module.’去掉，就是将torch.save(model.state_dict(), 'fff.pth'))改为torch.save(model.module.state_dict(), 'fff.pth'))

# 保存模型
model = torch.nn.DataParallel(model)
torch.save(model.module.state_dict(), 'fff.pth'))

# 加载模型
model.load_state_dict(torch.load(load_pth_name))

解决办法二

在加载的时候把参数名中的‘module.’去掉

# 保存模型
model = torch.nn.DataParallel(model)
torch.save(model.state_dict(), 'fff.pth'))

# 加载模型
state_dict = torch.load(load_pth_name)  
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:]  # remove `module.`
    new_state_dict[name] = v
model.load_state_dict(new_state_dict)

1.4 模型改变后参数加载

当使用像resnet50、resnet101这样的网络时，通常可以从网上对应下载到这些模型的预训练参数文件，但是当前所使用的模型，可能需要在resnet50或resnet101网络上进行一些修改，比如增加一些结构，或者删除一些结构。所以只希望加载修改后的模型与原来的模型之间具有相同结构部分的参数。

#第一种方法
#假设下载到的原有模型参数文件为checkpoint.pth.tar
model = OurModel()
model_checkpoint = torch.load('checkpoint.pth.tar')
pretrain_model_dict = model_checkpoint['state_dict']
model_dict = model.state_dict()
# 将两个模型参数进行比较，剔除不同的参数，保留相同的参数
same_model_dict = {k : v for k, v in pretrain_model_dict if k in model_dict}
# 更新相同的模型参数
model_dict.update(same_model_dict)
model.load_state_dict(model_dict)

#第二种方法
mymodelB = TheModelBClass(*args, **kwargs)
# strict=False，设置为false,只保留键值相同的参数
mymodelB.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)

# 加载我们真正需要的state_dict
mymodelB.load_state_dict(mymodelB_dict)

1.5 模型参数存储的内容

model.state_dict()

model.state_dict()返回的是一个OrderedDict对象。OrderedDict是dict的子类，其最大的特征就是可以维护添加的key-value对的顺序。说白了，OrderedDict也是一个字典，但是这个字典中的key-value对是有顺序的。这个特性正好可以跟网络结构的层次对应起来。

model.state.dict()的输出可以表示为以下形式，这里将输出中的tensor值简化为TensorValue

OrderedDict([('conv1.weight', TensorValue), ('conv1.bias', TensorValue),
             ('conv2.weight', TensorValue), ('conv2.bias', TensorValue)])

因此'conv1.weight'，'conv1.bias'，...等都是该字典中的键，键的名称如conv1、conv2都是在定义网络的时候自己设定的名称，而.weight和.bias都是固定表达。而TensorValue中则是值，其保存的都是网络中的参数值。

因此，我们可以通过遍历model.state_dict()的方式，获得各层参数值，但值得注意的是，只有那些参数可以训练的层，这些层的参数会被保存到model.state_dict()中，如卷积层、线性层，没有参数的层则不会被保存。

#遍历
for k, v in model.state_dict().items():
    print('k = ', k, '; ', 'v.size = ', v.size())   #为了直观，输出参数的尺寸大小

#输出
k =  conv1.weight ;  v.size =  torch.Size([3, 3, 3, 3])
k =  conv1.bias ;  v.size =  torch.Size([3])
k =  conv2.weight ;  v.size =  torch.Size([3, 3, 3, 3])
k =  conv2.bias ;  v.size =  torch.Size([3])

optimizer.state_dict()

不光模型有state_dict()，优化器也有该方法，它保存了优化器的状态以及被使用的超参数，使用方法如下：

optimizer = torch.optim.SGD(net.parameters(), lr=0.1, momentum=0.9, )

print(optimizer.state_dict())

#输出
{'state': {}, 'param_groups': [{'lr': 0.1, 'momentum': 0.9, 'dampening': 0, 
'weight_decay': 0, 'nesterov': False, 'params': [0, 1, 2, 3]}]}

参考链接：pytorch如何保存模型？ - 知乎

二、微调

用于特征提取的时候，要求特征提取部分参数不进行学习，而pytorch提供了requires_grad参数用于确定是否进去梯度计算，也即是否更新参数。以下以minist为例，用resnet18作特征提取：

#加载预训练模型
model = torchvision.models.resnet18(pretrained=True)

#遍历每一个参数，将其设置为不更新参数，即不学习
for param in model.parameters():
    param.requires_grad = False

# 将全连接层改为mnist所需的10类，注意：这样更改后requires_grad默认为True
model.fc = nn.Linear(512, 10)

# 优化
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)        

用于全局精调时，我们一般对不同的层需要设置不同的学习率，预训练的层学习率小一点，其他层大一点。

# 加载预训练模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 10)

# 参考：https://blog.csdn.net/u012759136/article/details/65634477
ignored_params = list(map(id, model.fc.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())

# 对不同参数设置不同的学习率
params_list = [{'params': base_params, 'lr': 0.001},]
params_list.append({'params': model.fc.parameters(), 'lr': 0.01})

optimizer = torch.optim.SGD(params_list,
                    0.001，
                    momentum=args.momentum,
                    weight_decay=args.weight_decay)

原文链接：PyTorch模型加载/参数初始化/Finetune - 知乎