support multi-gpu training

Author: hao-qiang

README.md

@@ -115,6 +115,11 @@ python train.py --dataset photo2cartoon
 python train.py --dataset photo2cartoon --pretrained_weights models/photo2cartoon_weights.pt
 ```
 
+多GPU训练(仍建议使用batch_size=1，单卡训练):
+```
+python train.py --dataset photo2cartoon --batch_size 4 --gpu_ids 0 1 2 3
+```
+
 ## Q&A
 #### Q：为什么开源的卡通化模型与小程序中的效果有差异？
 

README_EN.md

@@ -117,6 +117,11 @@ Load pre-trained weights:
 python train.py --dataset photo2cartoon --pretrained_weights models/photo2cartoon_weights.pt
 ```
 
+Train with Multi-GPU:
+```
+python train.py --dataset photo2cartoon --batch_size 4 --gpu_ids 0 1 2 3
+```
+
 ## Q&A
 #### Q：Why is the result of this project different from mini program?
 

models/UGATIT_sadalin_hourglass.py

@@ -44,7 +44,8 @@ def __init__(self, args):
         self.img_size = args.img_size
         self.img_ch = args.img_ch
 
-        self.device = args.device
+        self.device = f'cuda:{args.gpu_ids[0]}'
+        self.gpu_ids = args.gpu_ids
         self.benchmark_flag = args.benchmark_flag
         self.resume = args.resume
         self.rho_clipper = args.rho_clipper
@@ -154,6 +155,14 @@ def train(self):
             self.disLB.load_state_dict(params['disLB'])
             print(" [*] Load {} Success".format(self.pretrained_weights))
 
+        if len(self.gpu_ids) > 1:
+            self.genA2B = nn.DataParallel(self.genA2B, device_ids=self.gpu_ids)
+            self.genB2A = nn.DataParallel(self.genB2A, device_ids=self.gpu_ids)
+            self.disGA = nn.DataParallel(self.disGA, device_ids=self.gpu_ids)
+            self.disGB = nn.DataParallel(self.disGB, device_ids=self.gpu_ids)
+            self.disLA = nn.DataParallel(self.disLA, device_ids=self.gpu_ids)
+            self.disLB = nn.DataParallel(self.disLB, device_ids=self.gpu_ids)
+            
         # training loop
         print('training start !')
         start_time = time.time()
@@ -257,6 +266,9 @@ def train(self):
 
             G_id_loss_A = self.facenet.cosine_distance(real_A, fake_A2B)
             G_id_loss_B = self.facenet.cosine_distance(real_B, fake_B2A)
+            if len(self.gpu_ids) > 1:
+                G_id_loss_A = torch.mean(G_id_loss_A)
+                G_id_loss_B = torch.mean(G_id_loss_B)
 
             G_cam_loss_A = self.BCE_loss(fake_B2A_cam_logit, torch.ones_like(fake_B2A_cam_logit).to(self.device)) + \
                            self.BCE_loss(fake_A2A_cam_logit, torch.zeros_like(fake_A2A_cam_logit).to(self.device))
@@ -388,12 +400,22 @@ def train(self):
 
     def save(self, dir, step):
         params = {}
-        params['genA2B'] = self.genA2B.state_dict()
-        params['genB2A'] = self.genB2A.state_dict()
-        params['disGA'] = self.disGA.state_dict()
-        params['disGB'] = self.disGB.state_dict()
-        params['disLA'] = self.disLA.state_dict()
-        params['disLB'] = self.disLB.state_dict()
+        
+        if len(self.gpu_ids) > 1:
+            params['genA2B'] = self.genA2B.module.state_dict()
+            params['genB2A'] = self.genB2A.module.state_dict()
+            params['disGA'] = self.disGA.module.state_dict()
+            params['disGB'] = self.disGB.module.state_dict()
+            params['disLA'] = self.disLA.module.state_dict()
+            params['disLB'] = self.disLB.module.state_dict()            
+        
+        else:
+            params['genA2B'] = self.genA2B.state_dict()
+            params['genB2A'] = self.genB2A.state_dict()
+            params['disGA'] = self.disGA.state_dict()
+            params['disGB'] = self.disGB.state_dict()
+            params['disLA'] = self.disLA.state_dict()
+            params['disLB'] = self.disLB.state_dict()
         torch.save(params, os.path.join(dir, self.dataset + '_params_%07d.pt' % step))
 
     def load(self, dir, step):

train.py

@@ -31,7 +31,8 @@ def parse_args():
     parser.add_argument('--img_size', type=int, default=256, help='The size of image')
     parser.add_argument('--img_ch', type=int, default=3, help='The size of image channel')
 
-    parser.add_argument('--device', type=str, default='cuda:0', help='Set gpu mode: [cpu, cuda]')
+    # parser.add_argument('--device', type=str, default='cuda:0', help='Set gpu mode: [cpu, cuda]')
+    parser.add_argument('--gpu_ids', type=int, default=[0], nargs='+', help='Set [0, 1, 2, 3] for multi-gpu training')
     parser.add_argument('--benchmark_flag', type=str2bool, default=False)
     parser.add_argument('--resume', type=str2bool, default=False)
     parser.add_argument('--rho_clipper', type=float, default=1.0)