pytorch保存和加载模型权重以及CUDA在pytorch中的使用
1.保存加载自定义模型
1.1 保存加载整个模型
'''模型保存'''
torch.save(model, 'model.pkl')
'''模型加载'''
model = torch.load('model.pkl', map_location='cuda:0')
这种方式直接保存加载整个网络结构，比较死板，不能调整网络结构。
1.2 保存加载模型参数
'''模型参数保存'''
torch.save(model.state_dict(), 'model_param.pkl')
'''模型参数加载'''
device = torch.device('cpu') # device = torch.device('cuda')
# 定义网络
model = RNN() # 举例,一个RNN类,定义了RNN模型的结构
# 加载模型参数到模型结构
model.load_state_dict('model_param.pkl', map_location=device)
这种方式需要自己先定义网络模型的结构才能加载模型的参数，并且定义的网络模型的参数名称和结构要与加载的模型一致（可以是部分网络，比如只使用神经网络的前几层），相对灵活，便于对网络进行修改。
2.加载预训练模型
由于加载保存整个网络模型比较死板,所以一般都只保存或者加载预训练模型的参数.
2.1 预训练模型网络结构 == 自定义模型网络结构
结构相同就不需要修改,可以直接套用
path="预训练模型地址"
model = CJK_MODEL()
checkpoint = torch.load(path, map_location="cuda:0")
# 加载参数
model = model.load_state_dict(torch.load(path))
2.2 预训练模型网络结构与自定义模型网络结构不一致
- 首先打印出两个网络模型的各层网络名称
'''输出自定义模型的各层网络结构名称'''
model_dict = model.state_dict()
print(model_dict.keys())
'''输出自定义模型的各层网络结构名称'''
checkpoint = torch.load('./model_param.pkl')
for k, v in checkpoint.items():
print("keys:",k)
- 对比两者网络结构参数，如果差距太大就没有借用的必要了
- 如果许多参数层的名称完全一致:
model.load_state_dict(checkpoint, strict=True)
'''
load_state_dict 函数添加参数 strict=True，
它直接忽略那些没有的dict，有相同的就复制，没有就直接放弃赋值.
他要求预训练模型的关键字必须确切地严格地和
自定义网络的state_dict()函数返回的关键字相匹配才能赋值。
'''
- 如果许多参数层的名称大部分一致:
比如自定义网络模型中参数层名称为backbone.stage0.rbr_dense.conv.weight,
预训练模型中参数层名称为stage0.rbr_dense.conv.weight,可以看到二者大部分是一致的.这种情况下，可以把 预训练模型的stage0.rbr_dense.conv.weight读入网络的backbone.stage0.rbr_dense.conv.weight中。
3.cuda
把在cpu上运算的张量转变为在GPU上运行的cuda张量可以显著提升训练速度
- 在加载模型时转变为cuda张量
# 选择设备
device = ['cuda' if torch.cuda.is_available() else 'cpu']
# 法一:加载整个模型
model = torch.load('model.pkl', map_location='cuda:0')
# 法二:只加载模型参数
model = model.load_state_dict(torch.load(path), map_location='cuda:0')
- 单独转变为cuda张量
# 选择设备
device = ['cuda' if torch.cuda.is_available() else 'cpu']
# 法一:加载整个模型
model = torch.load('model.pkl')
# 法二:只加载模型参数
model = model.load_state_dict(torch.load(path))
'''单独转变为cuda张量'''
model.to(device)
注意, 必须把所有的张量都转为cuda型, 否则会报错