使用多块GPU 并行训练的代码

本文主要针对单主机多GPU的情形:
使用多块GPU所需的代码:
①在代码的最顶部使用如下代码:

device = 'cuda' if torch.cuda.is_available() else 'cpu'

注意:千万不要画蛇添足的写成 device = cuda:0 千万不要这样, 正确的做法是:去掉:0
接下来,将model 放入多块GPU:

    if torch.cuda.device_count() > 1:
        print("Let's use {0} GPUs!".format(torch.cuda.device_count()))
        print('可用的GPU数量:',torch.cuda.device_count())
        model = nn.DataParallel(model, device_ids=[0,2]).cuda()

②再使用.to(device)方法,将训练所需的图片、标签分别上传到GPU即可
如下所示:

 images = Variable(images).to(device)
            gts = Variable(gts).to(device)

你可能感兴趣的:(python,深度学习,人工智能)