pytorch多GPU并行计算及模型的保存与载入

1.多gpu并行计算

class Net(nn.Module):
    def __init__(input, output):
        pass
        #define your network 
net = Net(input, output) #实例化模型
net = nn.DataParallel(net) #数据并行
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #初始化计算设备
net .to(device)

2.模型的保存

if
torch.save(net.)

3.模型的重载

checkpoint = torch.load(resume)
state_dict =checkpoint['state_dict']

from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # remove 'module.' of dataparallel
    new_state_dict[name]=v

model.load_state_dict(new_state_dict)

4.模型的迁移

# cpu or gpu
torch.load('model/path', map_location='cpu')