pytorch-模型保存和加载

加载模型参数和选择是由保存的模型数据结构决定，故先要确定保存模型模型的方法和数据结构

保存模型

# 模型权重参数
model.state_dict()
‘‘‘首先说一下 model.state_dict()
pytorch 中的 model.state_dict 是一个简单的python的字典对象,将每一层与它的对应参数建立映射关系.(如model的每一层的weights及偏置等等)
只有那些参数可以训练的layer才会被保存到模型的state_dict中,如卷积层,线性层等

state_dict是在定义了model或optimizer之后pytorch自动生成的
‘‘‘
# model.state_dict() 其实返回的是一个OrderDict，存储了网络结构的名字和对应的参数


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.linear1 = nn.Linear(1, 2)
        self.linear2 = nn.Linear(2, 1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.linear2(x)

        return x


mode = Net()
print(mode.state_dict())
"""
OrderedDict([(‘linear1.weight‘, tensor([[ 0.8108],[-0.7968]])), (‘linear1.bias‘, tensor([ 0.2680, -0.4772])), (‘linear2.weight‘, tensor([[-0.7066, -0.3334]])), (‘linear2.bias‘, tensor([0.4819]))])

"""

print(mode.state_dict().keys())
"""
odict_keys([‘linear1.weight‘, ‘linear1.bias‘, ‘linear2.weight‘, ‘linear2.bias‘])
"""

for param_tensor in model.state_dict():
    #打印 key value字典
    print(param_tensor,‘\t‘,model.state_dict()[param_tensor].size())
    
"""
linear1.weight 	 torch.Size([2, 1])
linear1.bias 	 torch.Size([2])
linear2.weight 	 torch.Size([1, 2])
linear2.bias 	 torch.Size([1])
"""

# 保存模型

torch.save(obj, f, pickle_module,pickle_protocol )
"""输入参数
obj	   			可以是单个值也可以字典、对象
f 	   			要保存参数的文件路径
pickle_module
pickle_protocol
"""

# 1、自定义保存-工程实践中常常使用---推荐
state = {‘model‘:     model.state_dict(), 
         ‘optimizer‘: optimizer.state_dict(), 
         ‘epoch‘:     epoch   }
torch.save(model_object, ‘./model.pt‘)  

# 2、仅仅是保存模型权重参数
torch.save(model.state_dict(), PATH)

# 3、直接保存整个模型和模型结构
torch.save(Net,PATH)

加载模型

参数的保存

torch.save(model_object.state_dict(), ‘params.pth‘) 

# 模型的加载有模型保存的数据结构决定
ckpt = torch.load(f, map_location=None)
"""输入参数
f					file模型文件
map_location		torch.device, 动态地进行内存重映射，从不同的设备上读取文件
pickle_module 		用于unpickling元数据和对象的模块
pickle_load_args 	传递给pickle_module.load()

注释: 如果多块显卡，map_location={‘cuda:0‘:"cuda:1"},指定在2号显卡,不使用1号显卡
返回参数 字典d
由加载文件定义
默认情况，dict_keys([‘epoch‘, ‘state_dict‘, ‘optimizer‘, ‘best_pred‘])
"""

# 1、针对第一种保存模型的加载方式
# 加载模型
model=Net()										
# 加载模型参数
model_CKPT = torch.load(checkpoint_PATH) 
# 参数各个属性f
model.load_state_dict(model_CKPT[‘model‘])  
optimizer.load_state_dict(model_CKPT[‘optimizer‘])

# 2、针对第二种保存模型的加载方式
model=Net()									# 实例化网络
model_CKPT = torch.load(checkpoint_PATH)    # 加载模型参数
model.load_state_dict(model_CKPT)  

# 针对第三种保存整个模型的加载方式
model = torch.load(mode_PATH)

部分权重的加载

# 关键自定义函数

def intersect_dicts(da, db, exclude=()):
    """输入参数
    da (state_dict)			 加载权重的 state_dict 
    db (state_dict) 	 	 加载模型的 state_dict
    exclude (list)           不想要的权重 keys()
    
    返回参数
    加载的部分权重 (state_dict)
    """	
    ‘‘‘
    print("exclude",exclude)
    for k, v in da.items():
        for x in exclude:
            if x in k:
                print(‘@ ‘,x ,k)
            if v.shape != db[k].shape:
                print(‘# ‘, x, k)
	‘‘‘
    
    return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}

案例

# 加载模型
model = Net()

# 加载权重
ckpt=torch.load(weights, map_location=device)
state_dict=ckpt.state_dict()
# state_dict 是一个字典 

# state_dict.keys()
# odict_keys([‘0.model.0.conv.conv.weight‘, ‘0.model.0.conv.conv.bias‘, ‘0.model.1.conv.weight‘, .....])

# 权重取舍处理
state_dict=intersect_dicts(state_dict, model.state_dict(), exclude=exclude)

# 模型加载权重
model.load_state_dict(state_dict, strict=False)

# 最后可以输出加载了多少个
print(‘Transferred %g/%g items from %s‘ % (len(state_dict), len(model.state_dict()), weights))
# output >>> Transferred 498/506 items from yolov5m.pt

pytorch-模型保存和加载