train.py

import argparse
import os
import time
import torch
import numpy as np
import toml
from skimage.io import imsave
from encode.models import ENCODE
from encode.utils import (
    set_seed, print_cfg,
    worker_init_fn, CustomDataset, CustomDataset_4dstem,
    save_combined_checkpoint, count_parameters
)
from encode.utils import quant_utils
from encode.training import train, validate, validate_4dstem, train_quant
import warnings

from encode.utils.quant_utils import set_quantization


def _merge_asw_config(asw_toml):
    """Build ASW config from TOML section with defaults."""
    defaults = {
        'enabled': False,
        'alpha': 0.4,
        'update_freq': 5,
        'epoch_update_freq': 5,
        'num_samples': 64,
        'downsample_factor': 4,
        'spectral_index': 10,
        'replace_start': 0,
        'replace_end': 10,
        'only_decoder': True,
    }
    out = dict(defaults)
    for k, v in asw_toml.items():
        if k in out:
            out[k] = v
    return out


def load_config(config_path):
    """Load TOML config and convert to internal YAML-style format"""
    cfg_toml = toml.load(config_path)
    
    # Convert TOML to internal format
    model_cfg = cfg_toml['model']
    train_cfg = cfg_toml['training']
    quant_cfg = cfg_toml['quantization']
    data_cfg = cfg_toml['data']
    optim_cfg = cfg_toml['optimization']
    output_cfg = cfg_toml['output']
    
    # Build internal config dict (YAML-style)
    cfg = {
        'seed': train_cfg.get('seed', 42),
        'num_frames': data_cfg['num_frames'],
        'model': {
            'model_name': 'ENCODE',
            'embed_dim': model_cfg['embed_dim'],
            'fc_hw_dim': f"{model_cfg['fc_spatial'][0]}_{model_cfg['fc_spatial'][1]}_{model_cfg['fc_dim']}",
            'act': model_cfg['act'],
            'stride_list': model_cfg['stride_list'],
            't_dim': model_cfg['t_dim'],
            'expansion': model_cfg['expansion'],
            'reduction': model_cfg['reduction'],
            'lower_width': model_cfg['lower_width'],
            'bias': True,
            'norm': 'none',
            'conv_type': model_cfg['conv_type'],
            'out_act': model_cfg['out_act'],
            'out_channel': 1,
            'wavelet_levels': 1,
        },
        'epochs': train_cfg['epochs'],
        'batch_size': train_cfg['batch_size'],
        'num_workers': 4,
        'frame_gap': data_cfg['frame_gap'],
        'is_4dstem': False,
        'optim': {
            'optim_type': optim_cfg['optim'],
            'lr': train_cfg['lr'],
            'lr_min': train_cfg['lr_min'],
            'beta1': optim_cfg['beta1'],
            'beta2': optim_cfg['beta2'],
        },
        'loss': train_cfg['loss'],
        'warmup_epochs': train_cfg['warmup_epochs'],
        'grad_clip': train_cfg['grad_clip'],
        'accumulation_steps': train_cfg['accumulation_steps'],
        'scheduler': train_cfg['scheduler'],
        'quant_level': [quant_cfg['bits']],
        'quant_epochs': quant_cfg['qat_epochs'],
        'quant_noise': quant_cfg['noise_level'],
        'asw': _merge_asw_config(cfg_toml.get('asw', {})),
        'print_freq': output_cfg['print_freq'],
        'eval_freq': train_cfg['eval_freq'],
        'save_freq': output_cfg['save_freq'],
        'input_bit': data_cfg['input_bit'],
    }
    
    return cfg


def main(args):
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cuda.matmul.allow_tf32 = True

    cfg = load_config(args.config)

    if args.output_dir:
        output_dir = args.output_dir
    else:
        output_dir = cfg.get('output_dir', 'outputs')
    
    print_cfg(cfg)

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    save_dir = os.path.join(output_dir, 'images')
    os.makedirs(save_dir, exist_ok=True)

    checkpoints_dir = os.path.join(output_dir, 'checkpoints')
    os.makedirs(checkpoints_dir, exist_ok=True)

    bitstream_dir = os.path.join(output_dir, 'bitstream')
    os.makedirs(bitstream_dir, exist_ok=True)

    set_seed(cfg['seed'])

    device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
    print(f"\nUsing device: {device}\n")

    # Load data
    if cfg.get('is_4dstem', False):
        dataset = CustomDataset_4dstem(data_path=args.data, bit=cfg['input_bit'])
    else:
        dataset = CustomDataset(data_path=args.data, num_frames=cfg['num_frames'], frame_gap=cfg['frame_gap'], bit=cfg['input_bit'])
    data_size = dataset.data_size
    original_data_size_bytes = np.prod(data_size) * (cfg['input_bit'] // 8)
    original_data_size_mb = original_data_size_bytes / (1024 * 1024)
    num_frames = len(dataset)

    num_workers = max(cfg['num_workers'], 8)

    dataloader_train = torch.utils.data.DataLoader(
        dataset, batch_size=cfg['batch_size'], shuffle=True,
        num_workers=num_workers, pin_memory=True, worker_init_fn=worker_init_fn,
        persistent_workers=(num_workers > 0),
        prefetch_factor=4 if num_workers > 0 else None
    )
    dataloader_val = torch.utils.data.DataLoader(
        dataset, batch_size=1, shuffle=False, num_workers=num_workers,
        pin_memory=False, worker_init_fn=worker_init_fn,
        persistent_workers=(num_workers > 0)
    )

    model = ENCODE(cfg=cfg['model']).to(device)
    quant_utils.init_quantization(cfg, model)


    emb_h, emb_w, _ = [int(x) for x in cfg['model']['fc_hw_dim'].split('_')]
    _, emb_c = [int(x) for x in cfg['model']['embed_dim'].split('_')]

    params = count_parameters(model)
    embed_size = num_frames * (emb_h * emb_w * emb_c)
    model_size_mb = params / (1024*1024)
    embed_size_mb = embed_size / (1024*1024)
    total_size_mb = model_size_mb + embed_size_mb
    
    print("\n[Model Information]")
    print(f"  Model Parameters: {params/1e6:.3f}M")
    print(f"  Model Size:       {model_size_mb:.3f} MB")
    print(f"  Embedding Size:   {embed_size_mb:.3f} MB")
    print(f"  Total Size:       {total_size_mb:.3f} MB")
    print(f"  Original Data Size: {original_data_size_mb:.3f} MB")
    print(f"  Compression Ratio: {original_data_size_mb/total_size_mb:.2f}")
    print(f"  Number of Frames: {num_frames}")
    print()

    optim_cfg = cfg['optim']
    if optim_cfg['optim_type'] == 'Adam':
        optimizer = torch.optim.Adam(model.parameters(), lr=optim_cfg['lr'],
                                     betas=(optim_cfg['beta1'], optim_cfg['beta2']))
    else:
        raise NotImplementedError('optim_type: {} is not implemented.'.format(optim_cfg['optim_type']))

    # ACCELERATION: Use capturable scheduler
    scheduler_kwargs = {'T_max': cfg['epochs'], 'eta_min': cfg['optim']['lr_min']}
    if torch.cuda.is_available():
        import inspect
        sig = inspect.signature(torch.optim.lr_scheduler.CosineAnnealingLR.__init__)
        if 'capturable' in sig.parameters:
            scheduler_kwargs['capturable'] = True
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **scheduler_kwargs)


    qat_epochs = cfg['quant_epochs']
    
    # ACCELERATION: Initialize GradScaler for Mixed Precision (AMP)
    scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())

    # ASW settings
    use_asw = cfg.get('asw', {}).get('enabled', False)
    asw_cfg = cfg.get('asw', {})
    asw_state = {'S': None, 'hi': None, 'wi': None, 'last_update_epoch': -1}

    val_best_psnr, val_best_ssim = [np.array(0) for _ in range(2)]
    best_epoch = 0

    # Start training
    model.train()
    start_time = time.time()
    
    print("=" * 80)
    print("[Training Started]")
    print(f"  Total Epochs: {cfg['epochs']}")
    print(f"  Batch Size:   {cfg['batch_size']}")
    print(f"  Learning Rate: {cfg['optim']['lr']}")
    print(f"  Loss Type:    {cfg['loss']}")
    if cfg.get('warmup_epochs', 0) > 0:
        print(f"  Warmup:       {cfg['warmup_epochs']} epochs (linear 0 -> {cfg['optim']['lr']})")
    if cfg.get('accumulation_steps', 1) > 1:
        print(f"  Grad Accum:   {cfg['accumulation_steps']} steps (effective batch {cfg['accumulation_steps']})")
    if cfg.get('grad_clip', 0) > 0:
        print(f"  Grad Clip:    max_norm={cfg['grad_clip']}")
    if cfg.get('scheduler') == 'cosine':
        warmup = cfg.get('warmup_epochs', 0)
        print(f"  Scheduler:    warmup then cosine (T_max={cfg['epochs']-warmup})")
    print("=" * 80)
    print()

    for epoch in range(1, cfg['epochs'] + 1):
        # Training
        epoch_loss = train(dataloader_train, model, optimizer, scheduler, epoch, cfg['epochs'], cfg['loss'], device, cfg['print_freq'], scaler=scaler, use_asw=use_asw, asw_cfg=asw_cfg, asw_state=asw_state, grad_clip=cfg.get('grad_clip', 0))

        if epoch % cfg['eval_freq'] == 0 or epoch == cfg['epochs']:
            if cfg.get('is_4dstem', False):
                val_stats, embeddings_list = validate_4dstem(dataloader_val, model, epoch, device, save_dir, cfg['save_freq'], data_size, cfg['input_bit'], None)
            else:
                val_stats, embeddings_list = validate(dataloader_val, model, epoch, device, save_dir, cfg['save_freq'], num_frames, cfg['input_bit'])

            best_epoch = epoch if val_stats['val_psnr'] > val_best_psnr else best_epoch
            val_best_psnr = val_stats['val_psnr'] if val_stats['val_psnr'] > val_best_psnr else val_best_psnr
            val_best_ssim = val_stats['val_ssim'] if val_stats['val_ssim'] > val_best_ssim else val_best_ssim
            
            print("\n" + "-" * 80)
            print(f"[Validation Results - Epoch {epoch}]")
            print(f"  Current PSNR: {val_stats['val_psnr']:.2f} dB | SSIM: {val_stats['val_ssim']:.4f}")
            print(f"  Best PSNR:    {val_best_psnr:.2f} dB | SSIM: {val_best_ssim:.4f} (at Epoch {best_epoch})")
            print("-" * 80)
            
            if epoch == 1 or epoch == best_epoch:
                # Save combined checkpoint (model + embeddings) for best model
                best_checkpoint_path = os.path.join(checkpoints_dir, f'model_best.pth')
                save_combined_checkpoint(epoch, model, embeddings_list, best_checkpoint_path)


    print("\n" + "=" * 80)
    print("[Training Complete]")
    print(f"  Best PSNR: {val_best_psnr:.2f} dB (achieved at Epoch {best_epoch})")
    print(f"  Best SSIM: {val_best_ssim:.4f} (achieved at Epoch {best_epoch})")
    print("=" * 80)

    # QAT (Quantization-Aware Training)
    if qat_epochs > 0:
        print(f"\n{'=' * 40}")
        print(f"Start quant {cfg['quant_level'][0]}-bit fine-tuning for {qat_epochs} epochs")
        print(f"{'=' * 40}")

        q_level = cfg['quant_level'][0]
        q_noise = cfg['quant_noise']


        emb_full = torch.cat([item['emb'] for item in embeddings_list])
        emb_param = torch.nn.Parameter(emb_full.detach().clone().to(device), requires_grad=True)
        optimizer_qat = torch.optim.Adam([
            {'params': model.parameters()},
            {'params': emb_param}
        ], lr=1e-4)

        set_quantization(model, q_level, q_noise, False)

        for epoch in range(1, qat_epochs + 1):
            train_quant(dataloader_train, model, optimizer_qat, epoch, qat_epochs, cfg['loss'], device,
                               cfg['print_freq'], scaler, emb_param, q_level, num_frames)

        total_time = time.strftime("%H:%M:%S", time.gmtime(time.time() - start_time))
        print(f"  Total Training Time: {total_time}")

        num_bytes = quant_utils.compress_bitstream(model, bitstream_dir, q_level, emb_param)
        _, emb_recon = quant_utils.decompress(model, bitstream_dir, q_level)
        
        # Validate with quantized model
        val_stats, embeddings_list = validate(dataloader_val, model, epoch, device, save_dir, cfg['save_freq'],
                                              num_frames, cfg['input_bit'], emb_recon)
        print(f"[Validation Results - Epoch {epoch}]")
        print(f"  Current PSNR: {val_stats['val_psnr']:.2f} dB | SSIM: {val_stats['val_ssim']:.4f}")
        print(f'num_bytes: {num_bytes}')

        compressed_size_bytes = num_bytes
        compressed_size_mb = compressed_size_bytes / (1024 * 1024)

        compression_ratio = original_data_size_mb / compressed_size_mb if compressed_size_mb > 0 else 0
        space_saved_mb = original_data_size_mb - compressed_size_mb
        space_saved_percent = (space_saved_mb / original_data_size_mb * 100) if original_data_size_mb > 0 else 0

        print(f"  Original Data Size:    {original_data_size_mb:.3f} MB")
        print(f"  Compressed Size:       {compressed_size_mb:.3f} MB (model + embeddings)")
        print(f"  Compression Ratio:     {compression_ratio:.2f}:1")
        print(f"  Space Saved:           {space_saved_mb:.3f} MB ({space_saved_percent:.1f}%)")

        # Inference and save decompressed result
        model.eval()
        decompressed_frames = []
        input_bit = cfg.get('input_bit', 32)

        with torch.inference_mode():
            n_frames = emb_recon.shape[0]
            for valid_idx in range(n_frames):
                idx = valid_idx / n_frames
                emb = emb_recon[valid_idx, :].to(device, non_blocking=True)
                idx = torch.tensor(idx).to(device, non_blocking=True)
                with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
                    output, _ = model(None, idx, emb)
                output = output.detach().cpu()[0, 0].numpy()
                if input_bit == 16:
                    output = np.clip((output.astype(np.float32) * 65535), 0, 65535).astype(np.uint16)
                elif input_bit == 8:
                    output = np.clip((output.astype(np.float32) * 255), 0, 255).astype(np.uint8)
                decompressed_frames.append(output)

        decompressed_frames = np.stack(decompressed_frames, axis=0)
        decompressed_path = os.path.join(output_dir, 'decompressed_result.tif')
        imsave(decompressed_path, decompressed_frames, check_contrast=False)

        print(f"\n[Inference Complete]")
        print(f"  Decompressed result saved to: {decompressed_path}")
        print(f"  Output shape: {decompressed_frames.shape}")
        print("=" * 80)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='ENCODE Training')
    parser.add_argument('--config', type=str, required=True, help='Path to TOML config')
    parser.add_argument('--data', type=str, required=True, help='Path to input .tif file')
    parser.add_argument('--device', type=str, default='cuda:0', help='Device (e.g., cuda:0, cuda:1, cpu)')
    parser.add_argument('--output_dir', type=str, default=None, help='Output directory (overrides config)')
    args = parser.parse_args()
    
    main(args)