train2_stage2.py

import argparse
import logging
from logging import info as log
import os.path as path
import time
from datetime import datetime, timedelta
import yaml
import socket

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data.dataloader import DataLoader

import models.model as models
import utils.board as board
import utils.convert as convert
import utils.dataset as dataset
import utils.io as io
import utils.loss as loss
import utils.psnr as psnr
import utils.ssim as ssim

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
current_time = datetime.now().strftime("%Y-%m-%d_%H:%M:%S")
gaussian_blur = transforms.GaussianBlur(kernel_size=11, sigma=1.5)

# Directories initialization
checkpoint_dir = ""
debug = False

def init(config: dict):
    global checkpoint_dir
    global debug
    debug = config.get('debug') == True
    run_name = config['run_name']
    variation = config['variation']
    # Init directories
    output_dir = io.ensure_dir(path.join('output/train', run_name, variation, current_time))

    # Init tensorboard summary writer
    events_dir = io.ensure_dir(path.join(output_dir, 'events'))
    board.set_log_dir(events_dir)

    # Init stdout logging
    log_dir = io.ensure_dir(path.join(output_dir, 'logs'))
    logging.basicConfig(
        format="%(asctime)s | %(message)s",
        datefmt='%y-%m-%d %H:%M:%S',
        level=logging.DEBUG,
        handlers=[
            logging.FileHandler(path.join(log_dir, f"{current_time}.txt")),
            logging.StreamHandler()])

    checkpoint_dir = io.ensure_dir(path.join(output_dir, 'checkpoints'))
    if debug:
        log("Debug mode activated!")
    log(f"output: {output_dir}")
    log(f"host: {socket.gethostname()}")
    log(f"config:\n\n{yaml.dump(config, allow_unicode=True, default_flow_style=False, sort_keys=False)}")


def train(config, resume_checkpoint=None):
    # Init training
    best_epoch = 0
    best_val_loss = torch.inf
    best_metrics = []
    epoch_start = 0
    epoch_end = config['training']['epochs']

    # Init model
    model = models.ResHalfPredictor(
        train=True,
        stage=2,
        encoder_pretrained=config['model']["reshalf_pretrained"],
        invhalf_pretrained=config['model']["invhalf_pretrained"],
        use_input_y=config['model']['use_input_y'],
        noise_weight=config['model']['noise_weight']
        )

    # GPU enable
    model = nn.parallel.DataParallel(model).to(device)

    # Init optimizer
    optimizer: optim.Optimizer = getattr(
        optim, config['optimizer']['type'])(
        model.parameters(),
        **config['optimizer']['option'])
    log(f"Optimizer {optimizer.__class__.__name__} loaded.")

    # Init LR scheduler
    lr_scheduler = getattr(
        optim.lr_scheduler, config['lr_scheduler']['type'])(
        optimizer, **config['lr_scheduler']['option'])
    log(f"LR Scheduler {lr_scheduler.__class__.__name__} loaded.")

    # Load init checkpoint state
    init_checkpoint_path = config['model']['init_checkpoint']
    assert init_checkpoint_path is not None
    init_checkpoint = torch.load(init_checkpoint_path)
    model.load_state_dict(init_checkpoint['model_state_dict'], strict=False)
    log(f'Loaded init checkpoint {init_checkpoint_path}.')

    # Load checkpoint state
    if resume_checkpoint is not None and path.exists(resume_checkpoint):
        checkpoint = torch.load(resume_checkpoint)
        if checkpoint.get('epoch') is not None:
            epoch_start = checkpoint['epoch'] + 1
            log(f'Resume training from epoch {epoch_start}.')

        if checkpoint.get('model_state_dict') is not None:
            model.load_state_dict(checkpoint['model_state_dict'])
            log('Loaded model checkpoint.')

        if checkpoint.get("optimizer_state_dict") is not None:
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            log('Loaded optimizer checkpoint.')

        if checkpoint.get("lr_scheduler_state_dict") is not None:
            lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
            log('Loaded lr_sheduler checkpoint.')

        if checkpoint.get("best_val_loss") is not None:
            best_val_loss = checkpoint["best_val_loss"]
            log(f"Loaded best validation loss value {best_val_loss}")

        if checkpoint.get("best_metrics") is not None:
            best_metrics = checkpoint["best_metrics"]
            log('Loaded best validation metrics')

        if checkpoint.get("best_epoch") is not None:
            best_epoch = checkpoint["best_epoch"]
            log(f"Loaded best epoch {best_epoch}")

    # Init dataloaders
    dataset_root_dir = config['dataset']['root_dir']
    batch_size = config['dataset']['batch_size']
    num_workers = config['dataset']['num_workers']
    train_set = dataset.HalftoneVOC2012Training(root_dir=dataset_root_dir)
    train_data_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    log(f"Train DataLoader loaded with dataset: {train_set.__class__.__name__}")

    validate_set = dataset.HalftoneVOC2012Validation(root_dir=dataset_root_dir)
    validate_data_loader = DataLoader(validate_set, batch_size=batch_size, shuffle=False, num_workers=num_workers)
    log(f"Validate DataLoader loaded with dataset: {validate_set.__class__.__name__}")

    special_set = dataset.PlainColorTrainingDataset(root_dir=dataset_root_dir)
    special_data_loader = DataLoader(special_set, batch_size=batch_size, shuffle=True, num_workers=num_workers)
    log(f"Special DataLoader loaded with dataset: {special_set.__class__.__name__}")

    # Init loss functions
    loss_func = {
        "quantize_loss": loss.bin_l1,
        "tone_loss": loss.gaussian_l2,
        "structure_loss": loss.ssim_loss,
        "vgg_loss": loss.Vgg19Loss(),
        "blue_noise_loss": loss.l1_loss,
        "restore_loss": loss.l2_loss,
        "feature_loss": loss.l2_loss,
        "perceptual_loss": loss.l1_loss
    }

    # Run epochs
    epoch_count = 0
    metrics = []
    for epoch in range(epoch_start, epoch_end):
        if debug and epoch_count > 0:
            break
        log(f"Begin epoch {epoch}...")
        epoch_ts = time.time()
        epoch_lr = optimizer.state_dict()['param_groups'][0]['lr']

        train_loss = train_per_epoch(
            (epoch, epoch_end), config, model, optimizer=optimizer, train_loader=train_data_loader,
            special_loader=special_data_loader, loss_func=loss_func)
        val_loss, metrics = test_per_epoch(
            (epoch, epoch_end), config, model, test_data_loader=validate_data_loader, loss_func=loss_func)
        if lr_scheduler.__class__.__name__ == "ReduceLROnPlateau":
            lr_scheduler.step(val_loss)
        else:
            lr_scheduler.step()

        # Save checkpoint state
        if bool(val_loss < best_val_loss):
            best_val_loss = val_loss
            best_metrics = metrics
            best_epoch = epoch
        checkpoint_state = {
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'lr_scheduler_state_dict': lr_scheduler.state_dict(),
            'best_val_loss': best_val_loss,
            'best_metrics': best_metrics,
            'best_epoch': best_epoch
        }
        torch.save(checkpoint_state, path.join(checkpoint_dir, "model_latest.pth.tar"))
        if epoch == best_epoch:
            torch.save(checkpoint_state, path.join(checkpoint_dir, "model_best.pth.tar"))
            torch.save(checkpoint_state, path.join(checkpoint_dir, f"model_epoch{format(epoch, '03')}.pth.tar"))

        # Std log
        epoch_time_delta = time.time() - epoch_ts
        estimated_remaining_time = epoch_time_delta * (epoch_end - epoch - 1)
        log(f"End epoch {epoch}.")
        log(f"--- epoch:{epoch}/{epoch_end - 1}")
        log(f"--- train_loss: {round(train_loss, 6)}")
        log(f"--- val_loss: {round(val_loss, 6)}")
        log(f"--- best_epoch: {best_epoch}")
        log(f"--- time consumed: {timedelta(seconds=epoch_time_delta)} | estimated remaining time: {timedelta(seconds=estimated_remaining_time)}")

        # Log to tensorboard
        board.log_epoch_time(epoch, epoch_time_delta)
        board.log_epoch_lr(epoch, epoch_lr)
        board.log_epoch_metrics(epoch, metrics)
        epoch_count += 1

    log(f"All training iterations are done!")
    log(f"+++ Best epoch {best_epoch} / {epoch_end - 1}.")
    log(f"+++ Best val_loss {best_val_loss}.")
    log(f"+++ Best metrics:")
    board.log_epoch_metrics(best_epoch, best_metrics, to_tensorboard=False)
    # board.log_hparams(config, best_metrics)


def train_per_epoch(epoch: tuple, config: dict, model: nn.Module, optimizer: optim.Optimizer, train_loader: DataLoader,
                    special_loader: DataLoader, loss_func: dict):
    """Training iteration in a epoch

    Args:
        epoch (int, int): current epoch index and end epoch index
        config (dict): config dict
        model (nn.Module): model to be train
        optimizer (Optimizer): the training optimizer
        train_loader (DataLoader): dataloader for normal data
        special_loader (DataLoader): dataloader for special data
        loss_func (dict): dictionary contains loss functions
    """
    loss_weights = config['loss_weights']
    special_iter = iter(special_loader)
    total_batch = len(train_loader)
    total_loss = 0
    model.train()
    for batch_idx, (color_input, gray_input, cvt_input, halftone_input) in enumerate(train_loader):
        # We set this just for the example to run quickly.
        if debug and batch_idx > 0:
            break

        # Forward pass for normal data
        gray_input = gray_input.to(device)
        color_input = color_input.to(device)
        halftone_input = halftone_input.to(device)
        output = model(color_input, halftone_input)
        # halftone = output[0]
        # color_pred = output[1]
        y_pred = output[5]
        # crcb_pred = output[6]
        pseudo_y_pred = output[7]

        # Loss for normal data
        # if config['training']['ycrcb_supervised'] == True:
        #     restore_loss = loss_func['restore_loss'](crcb_pred, cvt_input[0][:, 1:, :, :].to(device))
        # else:
        #     restore_loss = loss_func['restore_loss'](color_pred, color_input)
        # feature_loss = loss_func['feature_loss'](pseudo_y_pred, gray_input)
        # vgg_loss = loss_func['vgg_loss'](convert.gray2bgr(y_pred / 2. + 0.5), convert.gray2bgr(gray_input / 2. + 0.5))
        # perceptual_loss = loss_func['perceptual_loss'](y_pred, gray_input)
        content_loss = loss_func['restore_loss'](pseudo_y_pred, gray_input)
        vgg_loss = loss_func['vgg_loss'](convert.gray2bgr(y_pred / 2. + 0.5), convert.gray2bgr(gray_input / 2. + 0.5))
        perceptual_loss = loss_func['perceptual_loss'](y_pred, gray_input)
        f_loss = loss_weights['vgg_loss_weight'] * vgg_loss + perceptual_loss

        loss = content_loss + loss_weights['perceptual_loss_weight'] * f_loss
        total_loss += loss.item()

        ##### Load special data ####
        # try:
        #     sp_input = next(special_iter)
        # except StopIteration:
        #     special_iter = iter(special_loader)
        #     sp_input = next(special_iter)

        # # Forward pass for special color
        # sp_color_input = sp_input[0].to(device)
        # sp_halftone_input = sp_input[3].to(device)
        # sp_output = model(sp_color_input, sp_halftone_input)
        # """
        # sp_output[0]: halftone
        # sp_output[1]: restored color of special color
        # sp_output[2]: DCT of halftone reference (ov)
        # sp_output[3]: DCT of halftone
        # sp_output[5]: y pred of halftone
        # sp_output[6]: crcb pred of halftone
        # """

        # # Loss for special color
        # if config['training']['ycrcb_supervised'] == True:
        #     _sp_restoreLoss = loss_func['restore_loss'](sp_output[6], sp_input[2][0][:, 1:, :, :].to(device))
        # else:
        #     _sp_restoreLoss = loss_func['restore_loss'](sp_output[1], sp_color_input)
        # _sp_feature_loss = loss_func['feature_loss'](sp_output[7], sp_input[1].to(device))

        # loss = loss_weights["restore_loss_weight"] * (restore_loss + 0 * _sp_restoreLoss) \
        #     + loss_weights["feature_loss_weight"] * (feature_loss + 0 * _sp_feature_loss) \
        #     + loss_weights["vgg_loss_weight"] * vgg_loss \
        #     + loss_weights["perceptual_loss_weight"] * perceptual_loss
        ##############################
        # loss = loss_weights["restore_loss_weight"] * restore_loss \
        #     + loss_weights["feature_loss_weight"] * feature_loss \
        #     + loss_weights["vgg_loss_weight"] * vgg_loss \
        #     + loss_weights["perceptual_loss_weight"] * perceptual_loss
        # total_loss += loss.item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (batch_idx + 1) % 100 == 0:
            log(f"[Train] >> epoch[{epoch[0]}/{epoch[1] - 1}] iter:{batch_idx + 1}/{total_batch} loss:{round(loss.item(), 6)}")

    return total_loss


def test_per_epoch(epoch: tuple, config: dict, model: nn.Module, test_data_loader: DataLoader, loss_func: dict):
    """Testing iteration in a epoch

    Args:
        epoch (int, int): current epoch index and end epoch index
        config (dict): config dict
        model (nn.Module): model to be inference
        test_data_loader (DataLoader): dataloader for normal data
        loss_func (dict): dictionary contains loss functions

    Returns:
        (float, [...]): the first value is the total loss, the second value is an array of metrics.
    """
    loss_weights = config['loss_weights']
    total_batch = len(test_data_loader)
    total_loss = 0
    all_metrics = [
        torch.zeros(total_batch, device=device, requires_grad=False),
        torch.zeros(total_batch, device=device, requires_grad=False),
        torch.zeros(total_batch, device=device, requires_grad=False),
        torch.zeros(total_batch, device=device, requires_grad=False)]
    model.eval()
    with torch.inference_mode():
        for batch_idx, (color_input, gray_input, cvt_input, halftone_input) in enumerate(test_data_loader):
            # We set this just for the example to run quickly.
            if debug and batch_idx > 0:
                break

            # Forward pass for normal data
            color_input = color_input.to(device)
            gray_input = gray_input.to(device)
            halftone_input = halftone_input.to(device)
            output = model(color_input, halftone_input)
            halftone = output[0]
            color_pred = output[1]
            halftone_q = output[4]
            y_pred = output[5]
            pseudo_y_pred = output[7]

            # Loss calculation
            content_loss = loss_func['restore_loss'](pseudo_y_pred, gray_input)
            vgg_loss = loss_func['vgg_loss'](convert.gray2bgr(y_pred / 2. + 0.5), convert.gray2bgr(gray_input / 2. + 0.5))
            perceptual_loss = loss_func['perceptual_loss'](y_pred, gray_input)
            f_loss = loss_weights['vgg_loss_weight'] * vgg_loss + perceptual_loss

            loss = content_loss + loss_weights['perceptual_loss_weight'] * f_loss
            total_loss += loss.item()

            # Metric calculation
            _color_input = convert.denormalize_tensor(color_input)
            _color_pred = convert.denormalize_tensor(color_pred)
            _halftone = convert.denormalize_tensor(halftone)
            _gray_input = convert.denormalize_tensor(gray_input)
            _halftone_q = convert.denormalize_tensor(halftone_q)
            _y_pred = convert.denormalize_tensor(y_pred)
            _blur_halftone = gaussian_blur(_halftone)
            _blur_gray_input = gaussian_blur(_gray_input)

            psnr_restore = psnr.psnr(_color_pred, _color_input)
            ssim_restore = ssim.ssim(_color_pred, _color_input)
            psnr_halftone = psnr.psnr(_blur_halftone, _blur_gray_input)
            ssim_halftone = ssim.ssim(_halftone, _gray_input)
            all_metrics[0][batch_idx] = psnr_restore
            all_metrics[1][batch_idx] = ssim_restore
            all_metrics[2][batch_idx] = psnr_halftone
            all_metrics[3][batch_idx] = ssim_halftone

            if (batch_idx + 1) % 100 == 0:
                log(f"[Test] >> epoch[{epoch[0]}/{epoch[1] - 1}] iter:{batch_idx + 1}/{total_batch} loss:{round(loss.item(), 6)}")

            if batch_idx == 0:
                board.log_images(epoch[0], 'color_input', convert.tensor_0_1(convert.bgr2rgb(_color_input)))
                board.log_images(epoch[0], 'color_output', convert.tensor_0_1(convert.bgr2rgb(_color_pred)))
                board.log_images(epoch[0], 'halftone', convert.tensor_0_1(_halftone))
                board.log_images(epoch[0], 'halftone_q', convert.tensor_0_1(_halftone_q))
                board.log_images(epoch[0], 'grayscale_output', convert.tensor_0_1(_y_pred))
                board.log_images(epoch[0], 'grayscale_input', convert.tensor_0_1(_gray_input))
    metrics = [
        # psnr_restore
        all_metrics[0].mean(),
        all_metrics[0].std(),
        # ssim_restore
        all_metrics[1].mean(),
        all_metrics[1].std(),
        # psnr_halftone
        all_metrics[2].mean(),
        all_metrics[2].std(),
        # ssim_halftone
        all_metrics[3].mean(),
        all_metrics[3].std(),
    ]

    return total_loss, metrics


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='train.py')
    parser.add_argument('-c', '--config', default=None, type=str, required=True)
    parser.add_argument('-r', '--resume', default=None, type=str, required=False)
    args = parser.parse_args()

    config: dict = {}
    with open(args.config) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)

    try:
        init(config)
        log("Program train2_stage2 start.")
        train(config, args.resume)
    except Exception:
        logging.error("Fatal Error in main loop", exc_info=True)
    finally:
        log("Program ended.")