nn.Dataparallel pytorch 平行计算的两种方法

1. nn.Dataparallel

多GPU加速训练

原理：
模型分别复制到每个卡中，然后把输入切片，分别放入每个卡中计算，然后再用第一块卡进行汇总求loss，反向传播更新参数。

第一块卡占用的内存多一点，因为output loss每次都会在第一块GPU相加计算，这就造成了第一块GPU的负载远远大于剩余其他的显卡。

要求：
batch_size > GPU 数量

第一种方法：

os.environment['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
device_ids = [0,1,2,3]
net  = torch.nn. Dataparallel(net, device_ids =device_ids)
net = net.cuda()

第二种方法

os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2"
if torch.cuda.is_available():
     self.device = "cuda"
     if torch.cuda.device_count() > 1:
         self.G = nn.DataParallel(self.G)
         self.D_A = nn.DataParallel(self.D_A)
         self.D_B = nn.DataParallel(self.D_B)
         self.vgg = nn.DataParallel(self.vgg)
         self.criterionHis = nn.DataParallel(self.criterionHis)
         self.criterionGAN = nn.DataParallel(self.criterionGAN)
         self.criterionL1 = nn.DataParallel(self.criterionL1)
         self.criterionL2 = nn.DataParallel(self.criterionL2)
         self.criterionGAN = nn.DataParallel(self.criterionGAN)
     self.G.cuda()
     self.vgg.cuda()
     self.criterionHis.cuda()
     self.criterionGAN.cuda()
     self.criterionL1.cuda()
     self.criterionL2.cuda()
     self.D_A.cuda()
     self.D_B.cuda()

2.模型分别单独放入每个指定的GPU中

把模型分别放到指定的GPU中，然后在运算的过程中，需要把利用**.to(cuda:x)** 去转移数据。这样暂用的内存比平行计算小。但是配置复杂一点。

	vgg_encoder = VGGEncoder().to('cuda:0')
	attn=CoAttention(channel=512).to('cuda:1')
	decoder = Decoder().to('cuda:2')
    optimizer_decoder = Adam(decoder.parameters(), lr=args.learning_rate)
    optimizer_attn = Adam(attn.parameters(), lr=args.learning_rate)
    
    content = content.cuda()  # 默认的是cuda:0
    style = style.cuda()
    content_features = vgg_encoder(content, output_last_feature=True)
    style_features = vgg_encoder(style, output_last_feature=True)
    content_features, style_features=attn(content_features.to('cuda:1'),style_features.to('cuda:1')) # 因为attn在cuda:1中