pytorch单机多gpu训练cycleGAN模型

废话不多说，直接上代码

修改cycleGan中的代码如下

原代码

disc_H = Discriminator(in_channels=3).to(config.DEVICE)disc_Z = Discriminator(in_channels=3).to(config.DEVICE)gen_Z = Generator(img_channels=3, num_residuals=9).to(config.DEVICE)gen_H = Generator(img_channels=3, num_residuals=9).to(config.DEVICE)

修改后的代码

   disc_H = Discriminator(in_channels=3)disc_Z = Discriminator(in_channels=3)gen_Z = Generator(img_channels=3, num_residuals=9)gen_H = Generator(img_channels=3, num_residuals=9)#使用多gpu加速if torch.cuda.device_count() > 1:print("Let's use", torch.cuda.device_count(), "GPUs!")disc_H = nn.DataParallel(disc_H, device_ids=[0,1])disc_Z = nn.DataParallel(disc_Z, device_ids=[0,1])gen_Z = nn.DataParallel(gen_Z, device_ids=[0,1])gen_H = nn.DataParallel(gen_H, device_ids=[0,1])disc_H.to(config.DEVICE)disc_Z.to(config.DEVICE)gen_Z.to(config.DEVICE)gen_H.to(config.DEVICE)

多gpu训练得到的结果单gpu貌似也能调用，效果如何，自行测试