From 2469bf15fea797074eece2e73cbaf69c21502296 Mon Sep 17 00:00:00 2001 From: budui Date: Thu, 3 Sep 2020 09:34:38 +0800 Subject: [PATCH] TAFG 0.01 --- configs/synthesizers/TAFG.yml | 9 +- configs/synthesizers/TAHG.yml | 129 ------------------ data/dataset.py | 28 +++- engine/TAFG.py | 8 +- engine/TAHG.py | 245 ---------------------------------- engine/base/i2i.py | 6 +- 6 files changed, 37 insertions(+), 388 deletions(-) delete mode 100644 configs/synthesizers/TAHG.yml delete mode 100644 engine/TAHG.py diff --git a/configs/synthesizers/TAFG.yml b/configs/synthesizers/TAFG.yml index f269560..b042ee4 100644 --- a/configs/synthesizers/TAFG.yml +++ b/configs/synthesizers/TAFG.yml @@ -49,7 +49,7 @@ loss: criterion: 'L1' style_loss: False perceptual_loss: True - weight: 1 + weight: 5 style: layer_weights: "1": 0.03125 @@ -63,10 +63,10 @@ loss: weight: 0 fm: level: 1 - weight: 1 + weight: 10 recon: level: 1 - weight: 1 + weight: 5 optimizers: generator: @@ -97,7 +97,8 @@ data: root_a: "/data/i2i/VoxCeleb2Anime/trainA" root_b: "/data/i2i/VoxCeleb2Anime/trainB" edges_path: "/data/i2i/VoxCeleb2Anime/edges" - edge_type: "hed" + landmarks_path: "/data/i2i/VoxCeleb2Anime/landmarks" + edge_type: "landmark_canny" size: [128, 128] random_pair: True pipeline: diff --git a/configs/synthesizers/TAHG.yml b/configs/synthesizers/TAHG.yml deleted file mode 100644 index 539b9a3..0000000 --- a/configs/synthesizers/TAHG.yml +++ /dev/null @@ -1,129 +0,0 @@ -name: TAHG -engine: TAHG -result_dir: ./result -max_pairs: 1000000 - -misc: - random_seed: 324 - -checkpoint: - epoch_interval: 1 # one checkpoint every 1 epoch - n_saved: 2 - -interval: - print_per_iteration: 10 # print once per 10 iteration - tensorboard: - scalar: 100 - image: 2 - -model: - generator: - _type: TAHG-Generator - _bn_to_sync_bn: False - style_in_channels: 3 - content_in_channels: 1 - num_blocks: 4 - discriminator: - _type: TAHG-Discriminator - in_channels: 3 - -loss: - gan: - loss_type: lsgan - real_label_val: 1.0 - fake_label_val: 0.0 - weight: 1.0 - edge: - criterion: 'L1' - hed_pretrained_model_path: "./network-bsds500.pytorch" - weight: 1 - perceptual: - layer_weights: - "3": 1.0 -# "0": 1.0 -# "5": 1.0 -# "10": 1.0 -# "19": 1.0 - criterion: 'L2' - style_loss: True - perceptual_loss: False - weight: 20 - recon: - level: 1 - weight: 1 - -optimizers: - generator: - _type: Adam - lr: 0.0001 - betas: [ 0.5, 0.999 ] - weight_decay: 0.0001 - discriminator: - _type: Adam - lr: 1e-4 - betas: [ 0.5, 0.999 ] - weight_decay: 0.0001 - -data: - train: - scheduler: - start_proportion: 0.5 - target_lr: 0 - buffer_size: 50 - dataloader: - batch_size: 160 - shuffle: True - num_workers: 2 - pin_memory: True - drop_last: True - dataset: - _type: GenerationUnpairedDatasetWithEdge - root_a: "/data/i2i/VoxCeleb2Anime/trainA" - root_b: "/data/i2i/VoxCeleb2Anime/trainB" - edges_path: "/data/i2i/VoxCeleb2Anime/edges" - edge_type: "hed" - size: [128, 128] - random_pair: True - pipeline: - - Load - - Resize: - size: [128, 128] - - ToTensor - - Normalize: - mean: [ 0.5, 0.5, 0.5 ] - std: [ 0.5, 0.5, 0.5 ] - test: - dataloader: - batch_size: 8 - shuffle: False - num_workers: 1 - pin_memory: False - drop_last: False - dataset: - _type: GenerationUnpairedDatasetWithEdge - root_a: "/data/i2i/VoxCeleb2Anime/testA" - root_b: "/data/i2i/VoxCeleb2Anime/testB" - edges_path: "/data/i2i/VoxCeleb2Anime/edges" - edge_type: "hed" - random_pair: False - size: [128, 128] - pipeline: - - Load - - Resize: - size: [128, 128] - - ToTensor - - Normalize: - mean: [ 0.5, 0.5, 0.5 ] - std: [ 0.5, 0.5, 0.5 ] - video_dataset: - _type: SingleFolderDataset - root: "/data/i2i/VoxCeleb2Anime/test_video_frames/" - with_path: True - pipeline: - - Load - - Resize: - size: [ 256, 256 ] - - ToTensor - - Normalize: - mean: [ 0.5, 0.5, 0.5 ] - std: [ 0.5, 0.5, 0.5 ] diff --git a/data/dataset.py b/data/dataset.py index fa7989b..6e62cdd 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -15,6 +15,7 @@ from tqdm import tqdm from .transform import transform_pipeline from .registry import DATASET +from .util import dlib_landmark def default_transform_way(transform, sample): @@ -178,20 +179,38 @@ class GenerationUnpairedDataset(Dataset): @DATASET.register_module() class GenerationUnpairedDatasetWithEdge(Dataset): - def __init__(self, root_a, root_b, random_pair, pipeline, edge_type, edges_path, size=(256, 256)): + def __init__(self, root_a, root_b, random_pair, pipeline, edge_type, edges_path, landmarks_path, size=(256, 256)): + assert edge_type in ["hed", "canny", "landmark_hed", "landmark_canny"] self.edge_type = edge_type self.size = size self.edges_path = Path(edges_path) + self.landmarks_path = Path(landmarks_path) assert self.edges_path.exists() + assert self.landmarks_path.exists() self.A = SingleFolderDataset(root_a, pipeline, with_path=True) self.B = SingleFolderDataset(root_b, pipeline, with_path=True) self.random_pair = random_pair def get_edge(self, origin_path): op = Path(origin_path) - edge_path = self.edges_path / f"{op.parent.name}/{op.stem}.{self.edge_type}.png" - img = Image.open(edge_path).resize(self.size) - return F.to_tensor(img) + if self.edge_type.startswith("landmark_"): + edge_type = self.edge_type.lstrip("landmark_") + use_landmark = True + else: + edge_type = self.edge_type + use_landmark = False + edge_path = self.edges_path / f"{op.parent.name}/{op.stem}.{edge_type}.png" + origin_edge = F.to_tensor(Image.open(edge_path).resize(self.size)) + if not use_landmark: + return origin_edge + else: + landmark_path = self.landmarks_path / f"{op.parent.name}/{op.stem}.{edge_type}.txt" + key_points, part_labels, part_edge = dlib_landmark.read_keypoints(landmark_path, size=self.size) + dist_tensor = torch.from_numpy(dlib_landmark.dist_tensor(key_points)) + part_labels = torch.from_numpy(part_labels) + edges = origin_edge * (part_labels.sum(0) == 0) # remove edges within face + edges = part_edge + edges + return torch.cat([edges, dist_tensor, part_labels], dim=0) def __getitem__(self, idx): a_idx = idx % len(self.A) @@ -200,7 +219,6 @@ class GenerationUnpairedDatasetWithEdge(Dataset): output["a"], path_a = self.A[a_idx] output["b"], path_b = self.B[b_idx] output["edge_a"] = self.get_edge(path_a) - output["edge_b"] = self.get_edge(path_b) return output def __len__(self): diff --git a/engine/TAFG.py b/engine/TAFG.py index be13eeb..72d3d89 100644 --- a/engine/TAFG.py +++ b/engine/TAFG.py @@ -65,12 +65,14 @@ class TAFGEngineKernel(EngineKernel): def criterion_generators(self, batch, generated) -> dict: loss = dict() - loss["perceptual"], _, = self.perceptual_loss(generated["b"], batch["b"]) * self.config.loss.perceptual.weight + loss_perceptual, _ = self.perceptual_loss(generated["b"], batch["a"]) + loss["perceptual"] = loss_perceptual * self.config.loss.perceptual.weight for phase in "ab": pred_fake = self.discriminators[phase](generated[phase]) - for i, sub_pred_fake in enumerate(pred_fake): + loss[f"gan_{phase}"] = 0 + for sub_pred_fake in pred_fake: # last output is actual prediction - loss[f"gan_{phase}_sub_{i}"] = self.gan_loss(sub_pred_fake[-1], True) + loss[f"gan_{phase}"] += self.gan_loss(sub_pred_fake[-1], True) if self.config.loss.fm.weight > 0 and phase == "b": pred_real = self.discriminators[phase](batch[phase]) diff --git a/engine/TAHG.py b/engine/TAHG.py deleted file mode 100644 index 71b705d..0000000 --- a/engine/TAHG.py +++ /dev/null @@ -1,245 +0,0 @@ -from itertools import chain -from math import ceil -from pathlib import Path - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torchvision - -import ignite.distributed as idist -from ignite.engine import Events, Engine -from ignite.metrics import RunningAverage -from ignite.utils import convert_tensor -from ignite.contrib.handlers.tensorboard_logger import OptimizerParamsHandler -from ignite.contrib.handlers.param_scheduler import PiecewiseLinear - -from omegaconf import OmegaConf, read_write - -import data -from loss.gan import GANLoss -from model.weight_init import generation_init_weights -from model.GAN.residual_generator import GANImageBuffer -from loss.I2I.edge_loss import EdgeLoss -from loss.I2I.perceptual_loss import PerceptualLoss -from util.image import make_2d_grid -from util.handler import setup_common_handlers, setup_tensorboard_handler -from util.build import build_model, build_optimizer - - -def build_lr_schedulers(optimizers, config): - g_milestones_values = [ - (0, config.optimizers.generator.lr), - (int(config.data.train.scheduler.start_proportion * config.max_iteration), config.optimizers.generator.lr), - (config.max_iteration, config.data.train.scheduler.target_lr) - ] - d_milestones_values = [ - (0, config.optimizers.discriminator.lr), - (int(config.data.train.scheduler.start_proportion * config.max_iteration), config.optimizers.discriminator.lr), - (config.max_iteration, config.data.train.scheduler.target_lr) - ] - return dict( - g=PiecewiseLinear(optimizers["g"], param_name="lr", milestones_values=g_milestones_values), - d=PiecewiseLinear(optimizers["d"], param_name="lr", milestones_values=d_milestones_values) - ) - - -def get_trainer(config, logger, train_data_loader): - generator = build_model(config.model.generator, config.distributed.model) - discriminators = dict( - a=build_model(config.model.discriminator, config.distributed.model), - b=build_model(config.model.discriminator, config.distributed.model), - ) - generation_init_weights(generator) - for m in discriminators.values(): - generation_init_weights(m) - - logger.debug(discriminators["a"]) - logger.debug(generator) - - optimizers = dict( - g=build_optimizer(generator.parameters(), config.optimizers.generator), - d=build_optimizer(chain(*[m.parameters() for m in discriminators.values()]), config.optimizers.discriminator), - ) - logger.info(f"build optimizers:\n{optimizers}") - - lr_schedulers = build_lr_schedulers(optimizers, config) - logger.info(f"build lr_schedulers:\n{lr_schedulers}") - - gan_loss_cfg = OmegaConf.to_container(config.loss.gan) - gan_loss_cfg.pop("weight") - gan_loss = GANLoss(**gan_loss_cfg).to(idist.device()) - - edge_loss_cfg = OmegaConf.to_container(config.loss.edge) - edge_loss_cfg.pop("weight") - edge_loss = EdgeLoss(**edge_loss_cfg).to(idist.device()) - - perceptual_loss_cfg = OmegaConf.to_container(config.loss.perceptual) - perceptual_loss_cfg.pop("weight") - perceptual_loss = PerceptualLoss(**perceptual_loss_cfg).to(idist.device()) - - recon_loss = nn.L1Loss() if config.loss.recon.level == 1 else nn.MSELoss() - - image_buffers = {k: GANImageBuffer(config.data.train.buffer_size or 50) for k in discriminators.keys()} - - def _step(engine, batch): - batch = convert_tensor(batch, idist.device()) - real = dict(a=batch["a"], b=batch["b"]) - fake = dict( - a=generator(content_img=batch["edge_a"], style_img=real["a"], which_decoder="a"), - b=generator(content_img=batch["edge_a"], style_img=real["b"], which_decoder="b"), - ) - rec_b = generator(content_img=batch["edge_b"], style_img=real["b"], which_decoder="b") - rec_bb = generator(content_img=batch["edge_b"], style_img=fake["b"], which_decoder="b") - - optimizers["g"].zero_grad() - loss_g = dict() - for d in "ab": - discriminators[d].requires_grad_(False) - pred_fake = discriminators[d](fake[d]) - loss_g[f"gan_{d}"] = config.loss.gan.weight * gan_loss(pred_fake, True) - _, t = perceptual_loss(fake[d], real[d]) - loss_g[f"perceptual_{d}"] = config.loss.perceptual.weight * t - loss_g[f"edge_{d}"] = config.loss.edge.weight * edge_loss(fake[d], batch["edge_a"]) - loss_g["recon_a"] = config.loss.recon.weight * recon_loss(fake["a"], real["a"]) - loss_g["recon_b"] = config.loss.recon.weight * recon_loss(rec_b, real["b"]) - loss_g["recon_bb"] = config.loss.recon.weight * recon_loss(rec_bb, real["b"]) - sum(loss_g.values()).backward() - optimizers["g"].step() - - for discriminator in discriminators.values(): - discriminator.requires_grad_(True) - - optimizers["d"].zero_grad() - loss_d = dict() - for k in discriminators.keys(): - pred_real = discriminators[k](real[k]) - pred_fake = discriminators[k](image_buffers[k].query(fake[k].detach())) - loss_d[f"gan_{k}"] = (gan_loss(pred_real, True, is_discriminator=True) + - gan_loss(pred_fake, False, is_discriminator=True)) / 2 - sum(loss_d.values()).backward() - optimizers["d"].step() - - generated_img = {f"real_{k}": real[k].detach() for k in real} - generated_img["rec_b"] = rec_b.detach() - generated_img["rec_bb"] = rec_b.detach() - generated_img.update({f"fake_{k}": fake[k].detach() for k in fake}) - generated_img.update({f"edge_{k}": batch[f"edge_{k}"].expand(-1, 3, -1, -1).detach() for k in "ab"}) - return { - "loss": { - "g": {ln: loss_g[ln].mean().item() for ln in loss_g}, - "d": {ln: loss_d[ln].mean().item() for ln in loss_d}, - }, - "img": generated_img - } - - trainer = Engine(_step) - trainer.logger = logger - for lr_shd in lr_schedulers.values(): - trainer.add_event_handler(Events.ITERATION_COMPLETED, lr_shd) - - RunningAverage(output_transform=lambda x: sum(x["loss"]["g"].values())).attach(trainer, "loss_g") - RunningAverage(output_transform=lambda x: sum(x["loss"]["d"].values())).attach(trainer, "loss_d") - - to_save = dict(trainer=trainer) - to_save.update({f"lr_scheduler_{k}": lr_schedulers[k] for k in lr_schedulers}) - to_save.update({f"optimizer_{k}": optimizers[k] for k in optimizers}) - to_save.update({"generator": generator}) - to_save.update({f"discriminator_{k}": discriminators[k] for k in discriminators}) - setup_common_handlers(trainer, config, to_save=to_save, clear_cuda_cache=True, set_epoch_for_dist_sampler=True, - end_event=Events.ITERATION_COMPLETED(once=config.max_iteration)) - - def output_transform(output): - loss = dict() - for tl in output["loss"]: - if isinstance(output["loss"][tl], dict): - for l in output["loss"][tl]: - loss[f"{tl}_{l}"] = output["loss"][tl][l] - else: - loss[tl] = output["loss"][tl] - return loss - - iter_per_epoch = len(train_data_loader) - tensorboard_handler = setup_tensorboard_handler(trainer, config, output_transform, iter_per_epoch) - if tensorboard_handler is not None: - tensorboard_handler.attach( - trainer, - log_handler=OptimizerParamsHandler(optimizers["g"], tag="optimizer_g"), - event_name=Events.ITERATION_STARTED(every=max(iter_per_epoch // config.interval.tensorboard.scalar, 1)) - ) - - @trainer.on(Events.ITERATION_COMPLETED(every=max(iter_per_epoch // config.interval.tensorboard.image, 1))) - def show_images(engine): - output = engine.state.output - image_order = dict( - a=["edge_a", "real_a", "fake_a", "fake_b"], - b=["edge_b", "real_b", "rec_b", "rec_bb"] - ) - for k in "ab": - tensorboard_handler.writer.add_image( - f"train/{k}", - make_2d_grid([output["img"][o] for o in image_order[k]]), - engine.state.iteration - ) - - with torch.no_grad(): - g = torch.Generator() - g.manual_seed(config.misc.random_seed) - random_start = torch.randperm(len(engine.state.test_dataset) - 11, generator=g).tolist()[0] - test_images = dict( - a=[[], [], [], []], - b=[[], [], [], []] - ) - for i in range(random_start, random_start + 10): - batch = convert_tensor(engine.state.test_dataset[i], idist.device()) - for k in batch: - batch[k] = batch[k].view(1, *batch[k].size()) - - real = dict(a=batch["a"], b=batch["b"]) - fake = dict( - a=generator(content_img=batch["edge_a"], style_img=real["a"], which_decoder="a"), - b=generator(content_img=batch["edge_a"], style_img=real["b"], which_decoder="b"), - ) - rec_b = generator(content_img=batch["edge_b"], style_img=real["b"], which_decoder="b") - rec_bb = generator(content_img=batch["edge_b"], style_img=fake["b"], which_decoder="b") - - test_images["a"][0].append(batch["edge_a"]) - test_images["a"][1].append(batch["a"]) - test_images["a"][2].append(fake["a"]) - test_images["a"][3].append(fake["b"]) - test_images["b"][0].append(batch["edge_b"]) - test_images["b"][1].append(batch["b"]) - test_images["b"][2].append(rec_b) - test_images["b"][3].append(rec_bb) - for n in "ab": - tensorboard_handler.writer.add_image( - f"test/{n}", - make_2d_grid([torch.cat(ti) for ti in test_images[n]]), - engine.state.iteration - ) - - return trainer - - -def run(task, config, logger): - assert torch.backends.cudnn.enabled - torch.backends.cudnn.benchmark = True - logger.info(f"start task {task}") - with read_write(config): - config.max_iteration = ceil(config.max_pairs / config.data.train.dataloader.batch_size) - - if task == "train": - train_dataset = data.DATASET.build_with(config.data.train.dataset) - logger.info(f"train with dataset:\n{train_dataset}") - train_data_loader = idist.auto_dataloader(train_dataset, **config.data.train.dataloader) - trainer = get_trainer(config, logger, train_data_loader) - if idist.get_rank() == 0: - test_dataset = data.DATASET.build_with(config.data.test.dataset) - trainer.state.test_dataset = test_dataset - try: - trainer.run(train_data_loader, max_epochs=ceil(config.max_iteration / len(train_data_loader))) - except Exception: - import traceback - print(traceback.format_exc()) - else: - return NotImplemented(f"invalid task: {task}") diff --git a/engine/base/i2i.py b/engine/base/i2i.py index 55ea67b..3592699 100644 --- a/engine/base/i2i.py +++ b/engine/base/i2i.py @@ -145,6 +145,7 @@ def get_trainer(config, ek: EngineKernel, iter_per_epoch): loss[tl] = output["loss"][tl] return loss + pairs_per_iteration = config.data.train.dataloader.batch_size tensorboard_handler = setup_tensorboard_handler(trainer, config, output_transform, iter_per_epoch) if tensorboard_handler is not None: tensorboard_handler.attach( @@ -159,7 +160,8 @@ def get_trainer(config, ek: EngineKernel, iter_per_epoch): test_images = {} for k in output["img"]: image_list = output["img"][k] - tensorboard_handler.writer.add_image(f"train/{k}", make_2d_grid(image_list), engine.state.iteration) + tensorboard_handler.writer.add_image(f"train/{k}", make_2d_grid(image_list), + engine.state.iteration * pairs_per_iteration) test_images[k] = [] for i in range(len(image_list)): test_images[k].append([]) @@ -182,6 +184,6 @@ def get_trainer(config, ek: EngineKernel, iter_per_epoch): tensorboard_handler.writer.add_image( f"test/{k}", make_2d_grid([torch.cat(ti) for ti in test_images[k]]), - engine.state.iteration + engine.state.iteration * pairs_per_iteration ) return trainer