Pytorch多GPU训练模型

一、指定一个GPU训练的两种方法：

1.代码中指定

import torch
torch.cuda.set_device(id)

2.终端中指定

CUDA_VISIBLE_DEVICES=1 python  你的程序

其中id就是你的gpu编号

二、多GPU并行训练：

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

该函数实现了在module级别上的数据并行使用，注意batch size要大于GPU的数量。

三、多GPU模型保存与加载

保存：

torch.save(model.module.state_dict(), model_out_path)

注意，要在模型名称的后面加上.module

加载：

from collections import OrderedDict
def myOwnLoad(model, check):
    modelState = model.state_dict()
    tempState = OrderedDict()
    for i in range(len(check.keys())-2):
        print modelState.keys()[i], check.keys()[i]
        tempState[modelState.keys()[i]] = check[check.keys()[i]]
    temp = [[0.02]*1024 for i in range(200)]  # mean=0, std=0.02
    tempState['myFc.weight'] = torch.normal(mean=0, std=torch.FloatTensor(temp)).cuda()
    tempState['myFc.bias']   = torch.normal(mean=0, std=torch.FloatTensor([0]*200)).cuda()

    model.load_state_dict(tempState)
    return model