使用多块GPU 并行训练的代码_使用两张cuda训练

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

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

注意：千万不要画蛇添足的写成 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()
1
2
3
4

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

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