lightning.py

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import pytorch_lightning as pl
import importlib
import PIL.Image as Image

import models
import datasets
from evaluator.ssim import SSIM, MSSSIM
import lpips
from models.loss import GANHingeLoss
from utils import set_logger, magic_image_handler

NUM_TEST_SAVE_IMAGE = 10


class FontLightningModule(pl.LightningModule):
    def __init__(self, args):
        super().__init__()
        self.args = args

        self.losses = {}
        self.metrics = {}
        self.networks = nn.ModuleDict(self.build_models())
        self.module_keys = list(self.networks.keys())

        self.losses = self.build_losses()
        self.metrics = self.build_metrics()

        self.opt_tag = {key: None for key in self.networks.keys()}
        self.sched_tag = {key: None for key in self.networks.keys()}
        self.sched_use = False
        # self.automatic_optimization = False

        self.train_d_content = True
        self.train_d_style = True

    def build_models(self):
        networks = {}
        for key, hp_model in self.args.models.items():
            key_ = key.lower()
            if 'g' == key_[0]:
                model_ = models.Generator(hp_model)
            elif 'd' == key_[0]:
                model_ = models.PatchGANDiscriminator(hp_model)  # TODO: add option for selecting discriminator
            else:
                raise ValueError(f"No key such as {key}")

            networks[key.lower()] = model_
        return networks

    def build_losses(self):
        losses_dict = {}
        losses_dict['L1'] = torch.nn.L1Loss()

        if 'd_content' in self.module_keys:
            losses_dict['GANLoss_content'] = GANHingeLoss()
        if 'd_style' in self.module_keys:
            losses_dict['GANLoss_style'] = GANHingeLoss()

        return losses_dict

    def build_metrics(self):
        metrics_dict = nn.ModuleDict()
        metrics_dict['ssim'] = SSIM(val_range=1)  # img value is in [0, 1]
        metrics_dict['msssim'] = MSSSIM(weights=[0.45, 0.3, 0.25], val_range=1)  # since imsize=64, len(weight)<=3
        metrics_dict['lpips'] = lpips.LPIPS(net='vgg')
        return metrics_dict

    def configure_optimizers(self):
        optims = {}
        for key, args_model in self.args.models.items():
            key = key.lower()
            if args_model['optim'] is not None:
                args_optim = args_model['optim']
                module, cls = args_optim['class'].rsplit(".", 1)
                O = getattr(importlib.import_module(module, package=None), cls)
                o = O([p for p in self.networks[key].parameters() if p.requires_grad],
                      lr=args_optim.lr, betas=args_optim.betas)

                optims[key] = o

        optim_module_keys = optims.keys()

        count = 0
        optim_list = []

        for _key in self.module_keys:
            if _key in optim_module_keys:
                optim_list.append(optims[_key])
                self.opt_tag[_key] = count
                count += 1

        return optim_list

    def forward(self, content_images, style_images):
        return self.networks['g']((content_images, style_images))

    def common_forward(self, batch, batch_idx):
        loss = {}
        logs = {}

        content_images = batch['content_images']
        style_images = batch['style_images']
        gt_images = batch['gt_images']
        image_paths = batch['image_paths']
        char_idx = batch['char_idx']

        generated_images = self(content_images, style_images)

        # l1 loss
        loss['g_L1'] = self.losses['L1'](generated_images, gt_images)
        loss['g_backward'] = loss['g_L1'] * self.args.logging.lambda_L1

        # loss for training generator
        if 'd_content' in self.module_keys:
            loss = self.d_content_loss_for_G(content_images, generated_images, loss)

        if 'd_style' in self.networks.keys():
            loss = self.d_style_loss_for_G(style_images, generated_images, loss)

        # loss for training discriminator
        generated_images = generated_images.detach()

        if 'd_content' in self.module_keys:
            if self.train_d_content:
                loss = self.d_content_loss_for_D(content_images, generated_images, gt_images, loss)

        if 'd_style' in self.module_keys:
            if self.train_d_style:
                loss = self.d_style_loss_for_D(style_images, generated_images, gt_images, loss)

        logs['content_images'] = content_images
        logs['style_images'] = style_images
        logs['gt_images'] = gt_images
        logs['generated_images'] = generated_images

        return loss, logs

    @property
    def automatic_optimization(self):
        return False

    def training_step(self, batch, batch_idx):
        metrics = {}
        # forward
        loss, logs = self.common_forward(batch, batch_idx)

        if self.global_step % self.args.logging.freq['train'] == 0:
            with torch.no_grad():
                metrics.update(self.calc_metrics(logs['gt_images'], logs['generated_images']))

        # backward
        opts = self.optimizers()

        opts[self.opt_tag['g']].zero_grad()
        self.manual_backward(loss['g_backward'])

        if 'd_content' in self.module_keys:
            if self.train_d_content:
                opts[self.opt_tag['d_content']].zero_grad()
                self.manual_backward(loss['dcontent_backward'])

        if 'd_style' in self.module_keys:
            if self.train_d_style:
                opts[self.opt_tag['d_style']].zero_grad()
                self.manual_backward(loss['dstyle_backward'])

        opts[self.opt_tag['g']].step()

        if 'd_content' in self.module_keys:
            if self.train_d_content:
                opts[self.opt_tag['d_content']].step()

        if 'd_style' in self.module_keys:
            if self.train_d_style:
                opts[self.opt_tag['d_style']].step()

        if self.global_step % self.args.logging.freq['train'] == 0:
            self.custom_log(loss, metrics, logs, mode='train')

    def validation_step(self, batch, batch_idx):
        metrics = {}
        loss, logs = self.common_forward(batch, batch_idx)
        self.custom_log(loss, metrics, logs, mode='eval')

    def test_step(self, batch, batch_idx):
        metrics = {}
        loss, logs = self.common_forward(batch, batch_idx)
        metrics.update(self.calc_metrics(logs['gt_images'], logs['generated_images']))

        if batch_idx < NUM_TEST_SAVE_IMAGE:
            for key, value in logs.items():
                if 'image' in key:
                    sample_images = (magic_image_handler(value) * 255)[..., 0].astype(np.uint8)
                    Image.fromarray(sample_images).save(f"{batch_idx:02d}_{key}.png")

        return loss, logs, metrics

    def test_epoch_end(self, test_step_outputs):
        # do something with the outputs of all test batches
        # all_test_preds = test_step_outputs.metrics
        ssim_list = []
        msssim_list = []

        for _, test_output in enumerate(test_step_outputs):

            ssim_list.append(test_output[2]['SSIM'].cpu().numpy())
            msssim_list.append(test_output[2]['MSSSIM'].cpu().numpy())

        print(f"SSIM: {np.mean(ssim_list)}")
        print(f"MSSSIM: {np.mean(msssim_list)}")

    def common_dataloader(self, mode='train', batch_size=None):
        dataset_cls = getattr(datasets, self.args.datasets.type)
        dataset_config = getattr(self.args.datasets, mode)
        dataset = dataset_cls(dataset_config, mode=mode)
        _batch_size = batch_size if batch_size is not None else dataset_config.batch_size
        dataloader = DataLoader(dataset,
                                shuffle=dataset_config.shuffle,
                                batch_size=_batch_size,
                                num_workers=dataset_config.num_workers,
                                drop_last=True)

        return dataloader

    def train_dataloader(self):
        return self.common_dataloader(mode='train')

    def val_dataloader(self):
        return self.common_dataloader(mode='eval')

    def test_dataloader(self):
        return self.common_dataloader(mode='eval')

    def calc_metrics(self, gt_images, generated_images):
        """

        :param gt_images:
        :param generated_images:
        :return:
        """
        metrics = {}
        _gt = torch.clamp(gt_images.clone(), 0, 1)
        _gen = torch.clamp(generated_images.clone(), 0, 1)
        metrics['SSIM'] = self.metrics['ssim'](_gt, _gen)
        msssim_value = self.metrics['msssim'](_gt, _gen)
        metrics['MSSSIM'] = msssim_value if not torch.isnan(msssim_value) else torch.tensor(0.).type_as(_gt)
        metrics['LPIPS'] = self.metrics['lpips'](_gt * 2 - 1, _gen * 2 - 1).squeeze().mean()
        return metrics

    # region step
    def d_content_loss_for_G(self, content_images, generated_images, loss):
        pred_generated = self.networks['d_content'](torch.cat([content_images, generated_images], dim=1))
        loss['g_gan_content'] = self.losses['GANLoss_content'](pred_generated, True, for_discriminator=False)

        loss['g_backward'] += loss['g_gan_content']
        return loss

    def d_content_loss_for_D(self, content_images, generated_images, gt_images, loss):
        # D
        if 'd_content' in self.module_keys:
            if self.train_d_content:
                pred_gt_images = self.networks['d_content'](torch.cat([content_images, gt_images], dim=1))
                pred_generated_images = self.networks['d_content'](torch.cat([content_images, generated_images], dim=1))

                loss['dcontent_gt'] = self.losses['GANLoss_content'](pred_gt_images, True, for_discriminator=True)
                loss['dcontent_gen'] = self.losses['GANLoss_content'](pred_generated_images, False, for_discriminator=True)
                loss['dcontent_backward'] = (loss['dcontent_gt'] + loss['dcontent_gen'])

        return loss

    def d_style_loss_for_G(self, style_images, generated_images, loss):
        pred_generated = self.networks['d_style'](torch.cat([style_images, generated_images], dim=1))
        loss['g_gan_style'] = self.losses['GANLoss_style'](pred_generated, True, for_discriminator=False)

        assert self.train_d_style
        loss['g_backward'] += loss['g_gan_style']
        return loss

    def d_style_loss_for_D(self, style_images, generated_images, gt_images, loss):
        pred_gt_images = self.networks['d_style'](torch.cat([style_images, gt_images], dim=1))
        pred_generated_images = self.networks['d_style'](torch.cat([style_images, generated_images], dim=1))

        loss['dstyle_gt'] = self.losses['GANLoss_style'](pred_gt_images, True, for_discriminator=True)
        loss['dstyle_gen'] = self.losses['GANLoss_style'](pred_generated_images, False, for_discriminator=True)
        loss['dstyle_backward'] = (loss['dstyle_gt'] + loss['dstyle_gen'])

        return loss

    def custom_log(self, loss, metrics, logs, mode):
        # logging values with tensorboard
        for loss_full_key, value in loss.items():
            model_type, loss_type = loss_full_key.split('_')[0], "_".join(loss_full_key.split('_')[1:])
            self.log(f'{model_type}/{mode}_{loss_type}', value)

        for metric_full_key, value in metrics.items():
            model_type, metric_type = metric_full_key.split('_')[0], "_".join(metric_full_key.split('_')[1:])
            self.log(f'{model_type}/{mode}_{metric_type}', value)

        # logging images, params, etc.
        tensorboard = self.logger.experiment
        for key, value in logs.items():
            if 'image' in key:
                sample_images = magic_image_handler(value)
                tensorboard.add_image(f"{mode}/" + key, sample_images, self.global_step, dataformats='HWC')
            elif 'param' in key:
                tensorboard.add_histogram(f"{mode}" + key, value, self.global_step)
            else:
                raise RuntimeError(f"Only logging with one of keywords: image, param | current input: {key}")