WaveDE.py

from __future__ import annotations

import argparse, math, os, re, csv
from pathlib import Path
from typing import List, Dict, Tuple, Any, Optional
import numpy as np
import torch
import torch.nn as nn
import torch.multiprocessing as mp
import torch.utils.checkpoint as checkpoint
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt
from tqdm import tqdm
import lpips
from datasets import PairFolder, PairFolderWithStem
import swanlab
from utils import run_validation, update_ema, save_checkpoint, setup_ddp, cleanup_ddp, is_main_process
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from loss import CharbonnierLoss, MSSSIMLoss, SemanticConsistencyLoss, FocalFrequencyLoss
from skmoe import FGMoE

os.environ["TOKENIZERS_PARALLELISM"] = "false"


def main():
    parser = argparse.ArgumentParser(description="SKCFG")
    parser.add_argument("--hq_root", required=True)
    parser.add_argument("--lq_root", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--test_hq_root", required=True)
    parser.add_argument("--test_lq_root", required=True)
    parser.add_argument("--resume_from", type=str, default=None)
    parser.add_argument("--steps", type=int, default=30000)
    parser.add_argument("--batch", type=int, default=32)
    parser.add_argument("--accum", type=int, default=1, help="Number of steps for gradient accumulation.")
    parser.add_argument("--res", type=int, default=256)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--bf16", action="store_true")
    parser.add_argument("--warmup_steps", type=int, default=1000)
    parser.add_argument("--model_id", type=str, required=True)
    parser.add_argument("--timestep", type=int, default=200)
    parser.add_argument("--gradient_checkpointing", action="store_true", help="Enable gradient checkpointing for UNet to save memory at the cost of speed.")
    parser.add_argument("--compile", action="store_true", help="Use torch.compile for UNet and VAE (requires PyTorch 2.0+).")
    parser.add_argument('--cfp_l1', type=float, default=0.7, help='Weight of Charbonnier L1 loss for CFP images.')
    parser.add_argument('--cfp_msssim', type=float, default=0.25, help='Weight of MS-SSIM loss for CFP images.')
    parser.add_argument("--path_vessel", type=str, required=True, help="Path to best_model.pth for ResNext50 Vessel model")
    parser.add_argument("--path_od", type=str, required=True, help="Path to PraNet-Best.pth for PraNet OD model")
    parser.add_argument("--lambda_vessel", type=float, default=0.01, help="Weight for Vessel semantic loss")
    parser.add_argument("--lambda_od", type=float, default=0.04, help="Weight for OD semantic loss")
    parser.add_argument("--lambda_ffl", type=float, default=0.1, help="Weight for Focal Frequency Loss")
    parser.add_argument("--ema", action="store_true")
    parser.add_argument("--ema_decay", type=float, default=0.999)
    parser.add_argument("--val_every", type=int, default=4000)
    parser.add_argument("--grad_clip_norm", type=float, default=1.0, help="Gradient clipping norm.")
    args = parser.parse_args()

    use_ddp = setup_ddp()
    if use_ddp:
        device = torch.device("cuda", int(os.environ["LOCAL_RANK"]))
        world_size = dist.get_world_size()
    else:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        world_size = 1

    torch.backends.cudnn.benchmark = True
    out_dir, ckpt_dir, val_dir = Path(args.out), Path(args.out)/"ckpts", Path(args.out)/"validation"
    log_dir = out_dir / "logs"

    writer = None
    if is_main_process():
        out_dir.mkdir(parents=True, exist_ok=True)
        ckpt_dir.mkdir(exist_ok=True)
        val_dir.mkdir(exist_ok=True)
        log_dir.mkdir(exist_ok=True)
        writer = SummaryWriter(log_dir=str(log_dir))
        
        swanlab.init(
            project="WaveDE",
            experiment_name="WaveDE",
            config=vars(args),
        )
    train_set = PairFolder(args.hq_root, args.lq_root, size=args.res, use_ram=False) 
    
    if use_ddp:
        train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=dist.get_rank(), shuffle=True)
        train_loader = DataLoader(train_set, batch_size=args.batch, sampler=train_sampler, num_workers=16, pin_memory=True, drop_last=True, persistent_workers=True)
    else:
        train_sampler = None
        train_loader = DataLoader(train_set, batch_size=args.batch, shuffle=True, num_workers=16, pin_memory=True, drop_last=True, persistent_workers=True, prefetch_factor=6)
    
    test_set = PairFolderWithStem(args.test_hq_root, args.test_lq_root, size=args.res, use_ram=False)
    
    if use_ddp:
        test_sampler = DistributedSampler(test_set, num_replicas=world_size, rank=dist.get_rank(), shuffle=False)
        test_loader = DataLoader(test_set, batch_size=args.batch * 3, sampler=test_sampler, num_workers=8)
    else:
        test_sampler = None
        test_loader = DataLoader(test_set, batch_size=args.batch * 3, shuffle=False, num_workers=8, pin_memory=True)

    G = FGMoE(args.timestep, args.model_id, use_gradient_checkpointing=args.gradient_checkpointing).to(device)
    
    if args.compile:
        G.unet = torch.compile(G.unet)
        G.vae = torch.compile(G.vae)
    
    if use_ddp:
        G = DDP(G, device_ids=[int(os.environ["LOCAL_RANK"])], find_unused_parameters=True)
        G_module = G.module
    else:
        G_module = G

    moe_params = []
    other_params = []
    moe_param_ids = {id(p) for p in G_module.moe_blocks.parameters()}
    
    for name, param in G_module.named_parameters():
        if not param.requires_grad:
            continue
        
        if id(param) in moe_param_ids:
            moe_params.append(param)
        else:
            other_params.append(param)
    optimizer_G = torch.optim.AdamW([
        {'params': moe_params, 'lr': args.lr * 2.0},
        {'params': other_params, 'lr': args.lr * 0.5}
    ], lr=args.lr)
    
    if is_main_process():
        def count_params(model):
            return sum(p.numel() for p in model.parameters() if p.requires_grad)

    loss_semantic_fn = SemanticConsistencyLoss(vessel_ckpt_path=args.path_vessel, od_ckpt_path=args.path_od, device=device).to(device)
    loss_l1_fn_mean = CharbonnierLoss(reduction='mean').to(device)
    loss_l1_fn_none = CharbonnierLoss(reduction='none').to(device)
    loss_msssim_fn_mean = MSSSIMLoss(reduction='mean').to(device)
    loss_msssim_fn_none = MSSSIMLoss(reduction='none').to(device)
    loss_ffl_fn = FocalFrequencyLoss(alpha=1.0, reduction='mean').to(device)
    weights_cfp = {'l1': args.cfp_l1, 'msssim': args.cfp_msssim}

    def lr_lambda(step):
        if step < args.warmup_steps:
            return float(step) / float(max(1, args.warmup_steps))
        progress = float(step - args.warmup_steps) / float(max(1, args.steps - args.warmup_steps))
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    scheduler_G = torch.optim.lr_scheduler.LambdaLR(optimizer_G, lr_lambda)

    ema_model = None
    if args.ema:
        ema_model = FGMoE(args.timestep, args.model_id).to("cpu").eval()
        state_dict = G_module.state_dict()
        new_state_dict = {}
        for k, v in state_dict.items():
            new_k = k.replace("_orig_mod.", "")
            new_state_dict[new_k] = v
            
        ema_model.load_state_dict(new_state_dict)

    step = 0
    best_val_loss = float('inf')

    if args.resume_from:
        ckpt = torch.load(args.resume_from, map_location=device)
        G_module.load_state_dict(ckpt['G'])
        optimizer_G.load_state_dict(ckpt['optimizer_G'])
        scheduler_G.load_state_dict(ckpt['scheduler_G'])
        step = ckpt.get('step', 0) + 1
        best_val_loss = ckpt.get('best_loss', float('inf'))
        if args.ema and 'ema' in ckpt:
            ema_model.load_state_dict(ckpt['ema'])

    train_iter = iter(train_loader)
    pbar = None
    if is_main_process():
        pbar = tqdm(initial=step, total=args.steps, desc="Training")
    last_best_ckpt_path = None
    
    try:
        while step < args.steps:
            for _ in range(args.accum):
                try:
                    lq, hq = next(train_iter)
                except StopIteration:
                    train_iter = iter(train_loader)
                    lq, hq = next(train_iter)
                
                lq, hq = lq.to(device), hq.to(device)

                with torch.amp.autocast(device_type='cuda', enabled=args.bf16, dtype=torch.bfloat16):
                    pred = G(lq)

                    l1_loss_map = loss_l1_fn_none(pred, hq)

                    with torch.amp.autocast("cuda", enabled=False):
                        pred_fp32 = pred.float()
                        hq_fp32 = hq.float()
                        loss_msssim_per_sample = loss_msssim_fn_none(torch.clamp(pred_fp32, -1, 1), torch.clamp(hq_fp32, -1, 1))
                    loss_msssim_per_sample = loss_msssim_per_sample.to(dtype=pred.dtype)
                    
                    l1_per_sample = l1_loss_map.mean(dim=[1,2,3])
                    loss_G_recons_per_sample = (l1_per_sample * weights_cfp['l1'] + loss_msssim_per_sample * weights_cfp['msssim'])
                        
                    loss_G_recons = loss_G_recons_per_sample.mean()

                    loss_vessel, loss_od = loss_semantic_fn(pred, hq)
                    
                    loss_ffl = loss_ffl_fn(pred_fp32, hq_fp32)

                    total_loss_G = (
                        loss_G_recons
                        + args.lambda_vessel * loss_vessel
                        + args.lambda_od * loss_od
                        + args.lambda_ffl * loss_ffl
                    )
                    
                    total_loss_G_accum = total_loss_G / args.accum
                
                total_loss_G_accum.backward()

            if args.grad_clip_norm > 0:
                torch.nn.utils.clip_grad_norm_(G_module.parameters(), args.grad_clip_norm)
            
            optimizer_G.step()
            scheduler_G.step()
            optimizer_G.zero_grad()

            if args.ema: update_ema(ema_model, G, args.ema_decay)

            if is_main_process():
                pbar.update(1)
                pbar.set_postfix({"Loss_G": f"{total_loss_G.item():.4f}", "LR": f"{scheduler_G.get_last_lr()[0]:.2e}"})
                
                if step % 10 == 0:
                    with torch.no_grad():
                        pred_for_logging = pred.float()
                        hq_for_logging = hq.float()
                        loss_l1_val = loss_l1_fn_mean(pred_for_logging, hq_for_logging)
                        loss_msssim_val = loss_msssim_fn_mean(torch.clamp(pred_for_logging, -1, 1), torch.clamp(hq_for_logging, -1, 1))

                    log_dict = {
                        "Loss/Generator_Total": total_loss_G.item(),
                        "Loss/Generator_Reconstruction": loss_G_recons.item(),
                        "Loss/L1_Charbonnier": loss_l1_val.item(),
                        "Loss/1-MSSSIM": loss_msssim_val.item(),
                        "Loss/Semantic_Vessel": loss_vessel.item(),
                        "Loss/Semantic_OD": loss_od.item(),
                        "Loss/FFL": loss_ffl.item(),
                        "Learning_Rate/Generator": scheduler_G.get_last_lr()[0],
                    }
                    
                    moe_log_dict = G_module.get_moe_logging_dict()
                    log_dict.update(moe_log_dict)
                    
                    swanlab.log(log_dict, step=step)

            if step > 0 and step % args.val_every == 0:
                if use_ddp:
                    dist.barrier()
                lpips_fn = lpips.LPIPS(net="vgg").to(device).eval()
                val_model = ema_model.to(device) if args.ema else G_module
                val_loss = run_validation(
                    model=val_model, test_loader=test_loader, 
                    loss_fns={'l1_none': loss_l1_fn_none, 'msssim_none': loss_msssim_fn_none},
                    weights=weights_cfp,
                    lpips_fn=lpips_fn, step=step, out_root=val_dir, device=device, writer=writer
                )
                if args.ema: val_model.to("cpu")

                if is_main_process() and val_loss < best_val_loss:
                    best_val_loss = val_loss

                    if last_best_ckpt_path and last_best_ckpt_path.exists():
                        last_best_ckpt_path.unlink()
                    
                    new_ckpt_path = ckpt_dir / f"best_model_step_{step}.pth"
                    last_best_ckpt_path = new_ckpt_path

                    raw_state = G_module.state_dict()
                    clean_state = {k.replace("_orig_mod.", ""): v for k, v in raw_state.items()}

                    state_dict = {
                        'G': clean_state, 
                        'optimizer_G': optimizer_G.state_dict(),
                        'scheduler_G': scheduler_G.state_dict(),
                        'step': step, 
                        'best_loss': best_val_loss
                    }
                    if args.ema: 
                        state_dict['ema'] = ema_model.state_dict()
                    
                    save_checkpoint(state_dict, new_ckpt_path)
                
                if use_ddp:
                    dist.barrier()
            
            step += 1
    
    finally:
        if is_main_process():
            if pbar: pbar.close()
            if writer: writer.close()
            swanlab.finish()
            final_path = ckpt_dir / "final_model.pth"
            raw_state = ema_model.state_dict() if args.ema else G_module.state_dict()
            final_model_state = {k.replace("_orig_mod.", ""): v for k, v in raw_state.items()}
            torch.save({'G': final_model_state}, final_path)
        
        cleanup_ddp()


if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    main()