pytorch CPU与GPU模型参数相互加载

---

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)