pytorch保存和加载模型权重以及CUDA在pytorch中的使用

1.保存加载自定义模型

1.1 保存加载整个模型

device = ['cuda:0' if torch.cuda.is_available() else 'cpu'][0]
'''模型保存'''
torch.save(model, 'model.pkl')
'''模型加载'''
model = torch.load('model.pkl', map_location=device)
1
2
3
4
5

这种方式直接保存加载整个网络结构，比较死板，不能调整网络结构。

1.2 保存加载模型参数

device = ['cuda:0' if torch.cuda.is_available() else 'cpu'][0]
# 定义网络
model = nn.RNN() # 举例,一个RNN类,定义了RNN模型的结构
'''模型参数保存'''
torch.save(model.state_dict(), 'model_param.pkl')

# 加载模型参数到模型结构
model.load_state_dict(torch.load('model_param.pkl', map_location=device))
1
2
3
4
5
6
7
8

这种方式需要自己先定义网络模型的结构才能加载模型的参数，并且定义的网络模型的参数名称和结构要与加载的模型一致（可以是部分网络，比如只使用神经网络的前几层），相对灵活，便于对网络进行修改。

2.加载预训练模型

由于加载保存整个网络模型比较死板,所以一般都只保存或者加载预训练模型的参数.

2.1 预训练模型网络结构 == 自定义模型网络结构

结构相同就不需要修改,可以直接套用

device = ['cuda:0' if torch.cuda.is_available() else 'cpu'][0]
path="预训练模型地址"
model = CJK_MODEL()
# 加载参数
model = model.load_state_dict(torch.load(path,mao_location=device))
1
2
3
4
5

2.2 预训练模型网络结构与自定义模型网络结构不一致

1. 首先打印出两个网络模型的各层网络名称

	device = ['cuda:0' if torch.cuda.is_available() else 'cpu'][0]
	'''输出自定义模型的各层网络结构名称'''
	model_dict = model.state_dict()
	print(model_dict.keys())
	'''输出自定义模型的各层网络结构名称'''
	checkpoint = torch.load('./model_param.pkl')
	for k, v in checkpoint.items():
		print("keys:",k)
1
2
3
4
5
6
7
8

2. 对比两者网络结构参数，如果差距太大就没有借用的必要了

• 如果许多参数层的名称完全一致:

	model.load_state_dict(checkpoint, strict=True)
	'''
	load_state_dict 函数添加参数 strict=True，
	它直接忽略那些没有的dict，有相同的就复制，没有就直接放弃赋值. 
	他要求预训练模型的关键字必须确切地严格地和
	自定义网络的state_dict()函数返回的关键字相匹配才能赋值。
	'''
1
2
3
4
5
6
7

• 如果许多参数层的名称大部分一致:
  比如自定义网络模型中参数层名称为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')
1
2
3
4
5
6

• 单独转变为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)
1
2
3
4
5
6
7
8

注意, 必须把所有的张量都转为cuda型, 否则会报错