1. 模型保存以及加载方法
# 直接保存模型 (参数 + 网络结构)
torch.save(model, '/path/to/save')
model = torch.load('/path/to/load')
# 只保存参数 (推荐)
torch.save(model.state_dict(), '/path/to/save')
model = NET()
mode.load_state_dict(torch.load('/path/to/load'))
# 保存参数、优化器、epoch
state = {
'model': model.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch
}
torch.save(state, '/path/to/save')
checkpoint = torch.load('/path/to/load')
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']
2. 单 GPU 和 单 CPU 参数-模型相互加载
## CPU->CPU OR GPU->GPU 直接加载
model.load_state_dict(torch.load('/path/to/load'))
## GPU->CPU (参数->模型)
state_dict = torch.load('/path/to/load', map_location=lambda storage, loc: storage)
state_dict = torch.load('/path/to/load', map_location='cpu')
model.load_state_dict(states_dict)
## CPU->GPU (参数->模型)
state_dict = torch.load('/path/to/load', map_location=lambda storage, loc: storage.cuda)
state_dict = torch.load('/path/to/load', map_location='cuda:0')
### 指定GPU
state_dict = torch.load('/path/to/load', map_location=lambda storage, loc: storage.cuda(1))
state_dict = torch.load('/path/to/load', map_location='cuda:1')
3. 多 GPU 模型-参数
## 模型 + 参数
torch.save(model.module, '/path/to/save') # 多了个module
## 参数
torch.save(model.module.state_dict(), '/path/to/save') # 多了个module
4. 单 GPU or CPU 模型加载多 GPU 参数
## 多gpu上保存的模型在参数名前多加了一个module.前缀
device = torch.device('cpu') # cup 模型
# device = torch.device('cuda:0') # gpu 模型
model = NET().to(device)
state_dict = torch.load('/path/to/load', map_location=device)
state_dict_new = {}
for k, v in state_dict.items():
new_k = k[7:] # 去掉键名的前七个字母,即'module.'
state_dict_new[new_k] = v
model.load_state_dict(state_dict_new)
5. 单 GPU or CPU 加载 多GPU模型+参数
model_cpu = NET().to('cpu')
model_gpu = NET().to('cuda:0')
pretrained_model = torch.load('/path/to/load') # 模型+参数
pretrained_dict = pretrained_model.module.state_dict() # 提取参数
model_cpu.load_state_dict(pretrained_dict)
model_gpu.load_state_dicr(pretrained_dicr)
6. 多 GPU 加载 多GPU参数
model = NET().to('cuda:0')
model = torch.nn.DataParallel(model, device_ids=[0, 1])
state_dict = torch.load('/path/to/load')
model.load_state_dict(state_dict)