红色大气企业网站,进度跟踪网站开发,知名广告公司,小马厂网站建设1. nn.Dataparallel
多GPU加速训练
原理#xff1a; 模型分别复制到每个卡中#xff0c;然后把输入切片#xff0c;分别放入每个卡中计算#xff0c;然后再用第一块卡进行汇总求loss#xff0c;反向传播更新参数。
第一块卡占用的内存多一点#xff0c;因为output los…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 cudaif 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)attnCoAttention(channel512).to(cuda:1)decoder Decoder().to(cuda:2)optimizer_decoder Adam(decoder.parameters(), lrargs.learning_rate)optimizer_attn Adam(attn.parameters(), lrargs.learning_rate)content content.cuda() # 默认的是cuda:0style style.cuda()content_features vgg_encoder(content, output_last_featureTrue)style_features vgg_encoder(style, output_last_featureTrue)content_features, style_featuresattn(content_features.to(cuda:1),style_features.to(cuda:1)) # 因为attn在cuda:1中