support dict dataset, clean code

configs/cyclegan-cityscapes.yaml

@@ -61,7 +61,7 @@ optimizer:
   beta1: 0.5
   lr_scheduler:
     name: linear
-    learning_rate: 0.0004
+    learning_rate: 0.0002
     start_epoch: 100
     decay_epochs: 100
 

ppgan/datasets/aligned_dataset.py

@@ -60,8 +60,8 @@ class AlignedDataset(BaseDataset):
         A = A_transform(A)
         B = B_transform(B)
 
-        # return {'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
-        return A, B, index #{'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
+        return {'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
+        # return A, B, index #{'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path}
 
     def __len__(self):
         """Return the total number of images in the dataset."""

ppgan/datasets/builder.py

 import paddle
 import numbers
 import numpy as np
+from multiprocessing import Manager
 from paddle.imperative import ParallelEnv
 
 from paddle.incubate.hapi.distributed import DistributedBatchSampler
@@ -10,25 +11,107 @@ from ..utils.registry import Registry
 DATASETS = Registry("DATASETS")
 
 
-def build_dataloader(cfg, is_train=True):
-    dataset = DATASETS.get(cfg.name)(cfg)
-    
-    batch_size = cfg.get('batch_size', 1)
+class DictDataset(paddle.io.Dataset):
+    def __init__(self, dataset):
+        self.dataset = dataset
+        self.tensor_keys_set = set()
+        self.non_tensor_keys_set = set()
+        self.non_tensor_dict = Manager().dict()
+
+        single_item = dataset[0]
+        self.keys = single_item.keys()
+        
+        for k, v in single_item.items():
+            if not isinstance(v, (numbers.Number, np.ndarray)):
+                self.non_tensor_dict.update({k: {}})
+                self.non_tensor_keys_set.add(k)
+            else:
+                self.tensor_keys_set.add(k)
+
+    def __getitem__(self, index):
+
+        ori_map = self.dataset[index]
+        
+        tmp_list = []
+        
+        for k, v in ori_map.items():
+            if isinstance(v, (numbers.Number, np.ndarray)):
+                tmp_list.append(v)
+            else:
+                tmp_dict = self.non_tensor_dict[k]
+                tmp_dict.update({index: v})
+                self.non_tensor_dict[k] = tmp_dict
 
-    # dataloader = DictDataLoader(dataset, batch_size, is_train)
+        tmp_list.append(index)
+        return tuple(tmp_list)
 
-    place = paddle.fluid.CUDAPlace(ParallelEnv().dev_id) \
+    def __len__(self):
+        return len(self.dataset)
+
+    def reset(self):
+        for k in self.non_tensor_keys_set:
+            self.non_tensor_dict[k] = {}
+
+
+class DictDataLoader():
+    def __init__(self, dataset, batch_size, is_train, num_workers=0):
+
+        self.dataset = DictDataset(dataset)
+
+        place = paddle.fluid.CUDAPlace(ParallelEnv().dev_id) \
                     if ParallelEnv().nranks > 1 else paddle.fluid.CUDAPlace(0)
 
-    sampler = DistributedBatchSampler(
-                dataset,
-                batch_size=batch_size,
-                shuffle=True if is_train else False,
-                drop_last=True if is_train else False)
+        sampler = DistributedBatchSampler(
+                                        self.dataset,
+                                        batch_size=batch_size,
+                                        shuffle=True if is_train else False,
+                                        drop_last=True if is_train else False)
+
+        self.dataloader = paddle.io.DataLoader(
+                                        self.dataset,
+                                        batch_sampler=sampler,
+                                        places=place,
+                                        num_workers=num_workers)
+
+        self.batch_size = batch_size
+
+    def __iter__(self):
+
+        self.dataset.reset()
+
+        for i, data in enumerate(self.dataloader):
+            return_dict = {}
+            j = 0
+            for k in self.dataset.keys:
+                if k in self.dataset.tensor_keys_set:
+                    return_dict[k] = data[j] if isinstance(data, (list, tuple)) else data
+                    j += 1
+                else:
+                    return_dict[k] = self.get_items_by_indexs(k, data[-1])
+            yield return_dict
+
+    def __len__(self):
+        return len(self.dataloader)
+
+    def get_items_by_indexs(self, key, indexs):
+        if isinstance(indexs, paddle.Variable):
+            indexs = indexs.numpy()
+        current_items = []
+        items = getattr(self.dataset, key)
+
+        for index in indexs:
+            current_items.append(items[index])
+
+        return current_items
+
+
+
+def build_dataloader(cfg, is_train=True):
+    dataset = DATASETS.get(cfg.name)(cfg)
+    
+    batch_size = cfg.get('batch_size', 1)
+    num_workers = cfg.get('num_workers', 0)
 
-    dataloader = paddle.io.DataLoader(dataset,
-                                      batch_sampler=sampler,
-                                      places=place,
-                                      num_workers=0)
+    dataloader = DictDataLoader(dataset, batch_size, is_train)
 
     return dataloader
\ No newline at end of file

ppgan/datasets/single_dataset.py

@@ -36,7 +36,13 @@ class SingleDataset(BaseDataset):
         # A_img = Image.open(A_path).convert('RGB')
         A_img = cv2.imread(A_path)
         A = self.transform(A_img)
-        return (A, index) #{'A': A, 'A_paths': A_path}
+        # items = {}
+        # if self.cfg.direction == 'AtoB':
+        #     items = {'A': A, 'A_paths': A_path}
+        # else:
+        #     items = {'B': A, 'B_paths': A_path}
+        # return items
+        return {'A': A, 'A_paths': A_path}
 
     def __len__(self):
         """Return the total number of images in the dataset."""

ppgan/datasets/unaligned_dataset.py

@@ -62,7 +62,8 @@ class UnalignedDataset(BaseDataset):
         A = self.transform_A(A_img)
         B = self.transform_B(B_img)
 
-        return A, B
+        # return A, B
+        return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}
 
     def __len__(self):
         """Return the total number of images in the dataset.

ppgan/engine/trainer.py

 import os
 import time
+
 import logging
 
 from paddle.imperative import ParallelEnv
@@ -7,7 +8,7 @@ from paddle.imperative import ParallelEnv
 from ..datasets.builder import build_dataloader
 from ..models.builder import build_model
 from ..utils.visual import tensor2img, save_image
-from ..utils.filesystems import save, load, makedirs
+from ..utils.filesystem import save, load, makedirs
 
 
 class Trainer:
@@ -45,6 +46,7 @@ class Trainer:
             for i, data in enumerate(self.train_dataloader):
                 self.batch_id = i
                 # unpack data from dataset and apply preprocessing
+                # data input should be dict
                 self.model.set_input(data)
                 self.model.optimize_parameters()
 
@@ -67,26 +69,21 @@ class Trainer:
         # test batch size must be 1
         for i, data in enumerate(self.test_dataloader):
             self.batch_id = i
-            # FIXME: dataloader not support map input, hard code now!!!
-            if self.cfg.dataset.test.name == 'AlignedDataset':
-                if self.cfg.dataset.test.direction == 'BtoA':
-                    fake = self.model.test(data[1])
-                else:
-                    fake = self.model.test(data[0])
-            elif self.cfg.dataset.test.name == 'SingleDataset':
-                fake = self.model.test(data[0])
-                
-            current_paths = self.test_dataloader.dataset.get_path_by_indexs(data[-1])
+
+            self.model.set_input(data)
+            self.model.test()
 
             visual_results = {}
+            current_paths = self.model.get_image_paths()
+            current_visuals = self.model.get_current_visuals()
+
             for j in range(len(current_paths)):
-                name = os.path.basename(current_paths[j])
-                name = os.path.splitext(name)[0]
+                short_path = os.path.basename(current_paths[j])
+                basename = os.path.splitext(short_path)[0]
+                for k, img_tensor in current_visuals.items():
+                    name = '%s_%s' % (basename, k)
+                    visual_results.update({name: img_tensor[j]})
 
-                visual_results.update({name + '_fakeB': fake[j]})
-            visual_results.update({name + '_realA': data[1]})
-            visual_results.update({name + '_realB': data[0]})
-            # visual_results.update({'realB': data[1]})
             self.visual('visual_test', visual_results=visual_results)
             
             if i % self.log_interval == 0:

ppgan/models/base_model.py

@@ -83,14 +83,14 @@ class BaseModel(ABC):
                 net = getattr(self, 'net' + name)
                 net.eval()
 
-    def test(self, input):
+    def test(self):
         """Forward function used in test time.
 
         This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
         It also calls <compute_visuals> to produce additional visualization results
         """
         with paddle.imperative.no_grad():
-            self.forward_test()
+            self.forward()
             self.compute_visuals()
 
     def compute_visuals(self):
@@ -105,7 +105,7 @@ class BaseModel(ABC):
         """Return visualization images. train.py will display these images with visdom, and save the images to a HTML"""
         visual_ret = OrderedDict()
         for name in self.visual_names:
-            if isinstance(name, str):
+            if isinstance(name, str) and hasattr(self, name):
                 visual_ret[name] = getattr(self, name)
         return visual_ret
 

ppgan/models/cycle_gan_model.py

@@ -84,32 +84,50 @@ class CycleGANModel(BaseModel):
 
         The option 'direction' can be used to swap domain A and domain B.
         """
-        AtoB = self.opt.dataset.train.direction == 'AtoB'
+        mode = 'train' if self.isTrain else 'test'
+        AtoB = self.opt.dataset[mode].direction == 'AtoB'
         
-        self.real_A = paddle.imperative.to_variable(input[0] if AtoB else input[1])
-        self.real_B = paddle.imperative.to_variable(input[1] if AtoB else input[0])
+        if AtoB:
+            if 'A' in input:
+                self.real_A = paddle.imperative.to_variable(input['A'])
+            if 'B' in input:
+                self.real_B = paddle.imperative.to_variable(input['B'])
+        else:
+            if 'B' in input:
+                self.real_A = paddle.imperative.to_variable(input['B'])
+            if 'A' in input:
+                self.real_B = paddle.imperative.to_variable(input['A'])
+
+        if 'A_paths' in input:
+            self.image_paths = input['A_paths']
+        elif 'B_paths' in input:
+            self.image_paths = input['B_paths']
+        # self.image_paths = input['A_paths' if AtoB else 'B_paths']
 
     def forward(self):
         """Run forward pass; called by both functions <optimize_parameters> and <test>."""
-        self.fake_B = self.netG_A(self.real_A)  # G_A(A)
-        self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
-        self.fake_A = self.netG_B(self.real_B)  # G_B(B)
-        self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
+        if hasattr(self, 'real_A'):
+            self.fake_B = self.netG_A(self.real_A)  # G_A(A)
+            self.rec_A = self.netG_B(self.fake_B)   # G_B(G_A(A))
 
+        if hasattr(self, 'real_B'):
+            self.fake_A = self.netG_B(self.real_B)  # G_B(B)
+            self.rec_B = self.netG_A(self.fake_A)   # G_A(G_B(B))
 
-    def forward_test(self, input):
-        input = paddle.imperative.to_variable(input)
-        net_g = getattr(self, 'netG_' + self.opt.dataset.test.direction[0])
-        return net_g(input)
 
-    def test(self, input):
-        """Forward function used in test time.
+    # def forward_test(self, input):
+    #     input = paddle.imperative.to_variable(input)
+    #     net_g = getattr(self, 'netG_' + self.opt.dataset.test.direction[0])
+    #     return net_g(input)
 
-        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
-        It also calls <compute_visuals> to produce additional visualization results
-        """
-        with paddle.imperative.no_grad():
-            return self.forward_test(input)
+    # def test(self, input):
+    #     """Forward function used in test time.
+
+    #     This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
+    #     It also calls <compute_visuals> to produce additional visualization results
+    #     """
+    #     with paddle.imperative.no_grad():
+    #         return self.forward_test(input)
 
     def backward_D_basic(self, netD, real, fake):
         """Calculate GAN loss for the discriminator

ppgan/models/pix2pix_model.py

@@ -83,8 +83,11 @@ class Pix2PixModel(BaseModel):
         # self.real_B = input['B' if AtoB else 'A'].to(self.device)
         # self.image_paths = input['A_paths' if AtoB else 'B_paths']
         AtoB = self.opt.dataset.train.direction == 'AtoB'
-        self.real_A = paddle.imperative.to_variable(input[0] if AtoB else input[1])
-        self.real_B = paddle.imperative.to_variable(input[1] if AtoB else input[0])
+        self.real_A = paddle.imperative.to_variable(input['A' if AtoB else 'B'])
+        self.real_B = paddle.imperative.to_variable(input['B' if AtoB else 'A'])
+        self.image_paths = input['A_paths' if AtoB else 'B_paths']
+        # self.real_A = paddle.imperative.to_variable(input[0] if AtoB else input[1])
+        # self.real_B = paddle.imperative.to_variable(input[1] if AtoB else input[0])
 
     def forward(self):
         """Run forward pass; called by both functions <optimize_parameters> and <test>."""
@@ -94,14 +97,14 @@ class Pix2PixModel(BaseModel):
         input = paddle.imperative.to_variable(input)
         return self.netG(input)
 
-    def test(self, input):
-        """Forward function used in test time.
+    # def test(self, input):
+    #     """Forward function used in test time.
 
-        This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
-        It also calls <compute_visuals> to produce additional visualization results
-        """
-        with paddle.imperative.no_grad():
-            return self.forward_test(input)
+    #     This function wraps <forward> function in no_grad() so we don't save intermediate steps for backprop
+    #     It also calls <compute_visuals> to produce additional visualization results
+    #     """
+    #     with paddle.imperative.no_grad():
+    #         return self.forward_test(input)
             
     def backward_D(self):
         """Calculate GAN loss for the discriminator"""

ppgan/utils/filesystems.py → ppgan/utils/filesystem.py

@@ -11,16 +11,15 @@ def save(state_dicts, file_name):
 
     def convert(state_dict):
         model_dict = {}
-        name_table = {}
+        # name_table = {}
         for k, v in state_dict.items():
             if isinstance(v, (paddle.framework.Variable, paddle.imperative.core.VarBase)):
                 model_dict[k] = v.numpy()
             else:
                 model_dict[k] = v
-                print('enter k', k)
                 return state_dict
-            name_table[k] = v.name
-        model_dict["StructuredToParameterName@@"] = name_table
+        #     name_table[k] = v.name
+        # model_dict["StructuredToParameterName@@"] = name_table
         return model_dict
 
     final_dict = {}