diff --git a/configs/cyclegan-cityscapes.yaml b/configs/cyclegan-cityscapes.yaml index fb61423dcea60b397c1de03c53a4492e528d5cc5..fd48f9f5905967e11bf4db336a4dbef23d0cd49f 100644 --- a/configs/cyclegan-cityscapes.yaml +++ b/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 diff --git a/ppgan/datasets/aligned_dataset.py b/ppgan/datasets/aligned_dataset.py index d10583e4ccfba2e67c83de367a0e31ebea9662f4..63db1c75c515ea70e968193cb08db7667e8cbd71 100644 --- a/ppgan/datasets/aligned_dataset.py +++ b/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.""" diff --git a/ppgan/datasets/builder.py b/ppgan/datasets/builder.py index 67accf2395aa8f1b68399ece1a0ad63f4267e61e..e94cebafc5884f0921b3e2f8dffa6146a5350172 100644 --- a/ppgan/datasets/builder.py +++ b/ppgan/datasets/builder.py @@ -1,6 +1,7 @@ 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 diff --git a/ppgan/datasets/single_dataset.py b/ppgan/datasets/single_dataset.py index 1fa76d253f4e13794209dbffc62ca6184519898c..97e91d5d06e4cafb6b2057ed5953b09a01131b46 100644 --- a/ppgan/datasets/single_dataset.py +++ b/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.""" diff --git a/ppgan/datasets/unaligned_dataset.py b/ppgan/datasets/unaligned_dataset.py index 6e95e519172bcf6f90a81b09ad54c0d94adf4169..da673a07216cf6ae827b5431d6fd42f86ca1a0dc 100644 --- a/ppgan/datasets/unaligned_dataset.py +++ b/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. diff --git a/ppgan/engine/trainer.py b/ppgan/engine/trainer.py index 45bc58b1e91815ce4c06e677ace2e35b541eb485..35dc5fd8a988148367969ce12511ba5c6eec0e43 100644 --- a/ppgan/engine/trainer.py +++ b/ppgan/engine/trainer.py @@ -1,5 +1,6 @@ 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: diff --git a/ppgan/models/base_model.py b/ppgan/models/base_model.py index f673b899f6a7549398c24644800a4b4270896679..6d038c77008e32002eff09c46f06f474567c93a9 100644 --- a/ppgan/models/base_model.py +++ b/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 function in no_grad() so we don't save intermediate steps for backprop It also calls 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 diff --git a/ppgan/models/cycle_gan_model.py b/ppgan/models/cycle_gan_model.py index f4564e7bcfe80852fa15bd51bd0d6c8ee581d2a1..2d4b9e3998a5cce832aa54911ace28699193c2d5 100644 --- a/ppgan/models/cycle_gan_model.py +++ b/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 and .""" - 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 function in no_grad() so we don't save intermediate steps for backprop - It also calls 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 function in no_grad() so we don't save intermediate steps for backprop + # It also calls 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 diff --git a/ppgan/models/pix2pix_model.py b/ppgan/models/pix2pix_model.py index cd14efce32a7c8e4ca25448f9ca830b592f42646..ab3ff48bd05808bd3259339b21c18679ae989057 100644 --- a/ppgan/models/pix2pix_model.py +++ b/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 and .""" @@ -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 function in no_grad() so we don't save intermediate steps for backprop - It also calls to produce additional visualization results - """ - with paddle.imperative.no_grad(): - return self.forward_test(input) + # This function wraps function in no_grad() so we don't save intermediate steps for backprop + # It also calls 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""" diff --git a/ppgan/utils/filesystems.py b/ppgan/utils/filesystem.py similarity index 87% rename from ppgan/utils/filesystems.py rename to ppgan/utils/filesystem.py index 2b058c98651e1f91b2282605d6fae7b1700f9ec3..1eb9f0da435fe8eb945025717f1d591ea2cd9c71 100644 --- a/ppgan/utils/filesystems.py +++ b/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 = {}