finetune.py

"""
GraPHFormer Fine-tuning Script

Fine-tune pretrained CLIP-style models for classification.
Supports: image_only, tree_only, multimodal modes.

Usage:
    python finetune.py --exp_name my_finetune --pretrained_checkpoint path/to/checkpoint.pth
"""

import argparse
import datetime
import time
import os
import json
import numpy as np

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from sklearn.neighbors import KNeighborsClassifier

from graphformer.models import CLIPModel, FineTuneModel
from graphformer.augmentations import (
    Compose,
    RandomScaleCoords, RandomRotate, RandomJitter, RandomShift,
    RandomFlip, RandomMaskFeats, RandomJitterLength, RandomElasticate,
    RandomDropSubTrees, RandomSkipParentNode, RandomSwapSiblingSubTrees,
    CombinedPersistenceAugmentation,
)
from graphformer.utils import save_checkpoint, get_root_logger, set_seed
from graphformer.data import NeuronTreeDataset, get_collate_fn, LABEL_DICT


def evaluate_accuracy(model, data_loader, device):
    """Evaluate classification accuracy"""
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for batch in data_loader:
            batch = batch.to(device)
            _, logits = model(batch)

            pred = logits.argmax(dim=1)
            labels = batch.label.cuda() if not batch.label.is_cuda else batch.label

            correct += (pred == labels).sum().item()
            total += labels.size(0)

    accuracy = correct / total * 100
    model.train()
    return accuracy


def extract_features(model, data_loader, device):
    """Extract features for KNN evaluation"""
    model.eval()
    features_list = []
    labels_list = []

    with torch.no_grad():
        for batch in data_loader:
            batch = batch.to(device)
            _, _, features = model(batch, return_features=True)

            features_list.append(features)
            labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
            labels_list.append(labels)

    features = torch.cat(features_list, dim=0)
    labels = torch.cat(labels_list, dim=0)

    model.train()
    return features, labels


def evaluate_knn(model, train_loader, test_loader, device, knn_k=20):
    """KNN evaluation using sklearn"""
    x_train, y_train = extract_features(model, train_loader, device)
    x_test, y_test = extract_features(model, test_loader, device)

    neigh = KNeighborsClassifier(n_neighbors=knn_k)
    neigh.fit(x_train.cpu().numpy(), y_train.cpu().numpy())

    score = neigh.score(x_test.cpu().numpy(), y_test.cpu().numpy())

    return score * 100


def mixup_data(features, labels, alpha=1.0):
    """Apply mixup augmentation"""
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = features.size(0)
    index = torch.randperm(batch_size).to(features.device)

    mixed_features = lam * features + (1 - lam) * features[index, :]
    labels_a, labels_b = labels, labels[index]

    return mixed_features, labels_a, labels_b, lam


def mixup_criterion(criterion, pred, labels_a, labels_b, lam):
    """Mixup loss"""
    return lam * criterion(pred, labels_a) + (1 - lam) * criterion(pred, labels_b)


if __name__ == "__main__":
    parser = argparse.ArgumentParser("GraPHFormer Fine-tuning")

    # Basic
    parser.add_argument("--work_dir", type=str, default="./work_dir")
    parser.add_argument("--exp_name", type=str, required=True)
    parser.add_argument("--dataset", type=str, default="bil_6_classes")
    parser.add_argument("--data_dir", type=str, default="data/raw/bil")
    parser.add_argument("--seed", type=int, default=42)

    # Pretrained checkpoint
    parser.add_argument("--pretrained_checkpoint", type=str, default=None)

    # Fine-tuning mode
    parser.add_argument("--mode", type=str, default="multimodal",
                        choices=["image_only", "tree_only", "multimodal"])
    parser.add_argument("--freeze_encoders", action="store_true", default=False)
    parser.add_argument("--freeze_image_only", action="store_true", default=False)
    parser.add_argument("--linear_probe_epochs", type=int, default=0)
    parser.add_argument("--use_projection", action="store_true", default=False)
    parser.add_argument("--fusion_mode", type=str, default="concat",
                        choices=["concat", "add", "cross_attention", "bi_attention", "gated", "cmf", "mhcma"])

    # Tree Model
    parser.add_argument("--tree_model", type=str, default="double",
                        choices=["ori", "v2", "double"])
    parser.add_argument("--child_mode", type=str, default="sum")
    parser.add_argument("--input_features", nargs="+", type=int,
                        default=[2, 3, 4, 12, 13])
    parser.add_argument("--h_size", type=int, default=256)
    parser.add_argument("--bn", action="store_true", default=False)

    # Image Model
    parser.add_argument("--image_encoder", type=str, default="resnet18")
    parser.add_argument("--image_size", type=int, default=256)
    parser.add_argument("--freeze_image_backbone", action="store_true", default=False)

    # CLIP settings
    parser.add_argument("--embed_dim", type=int, default=128)
    parser.add_argument("--temperature", type=float, default=0.07)
    parser.add_argument("--loss_type", type=str, default="clip")

    # Training
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--epochs", default=50, type=int)
    parser.add_argument("--lr", default=1e-4, type=float)
    parser.add_argument("--wd", default=0.01, type=float)
    parser.add_argument("--warmup_epochs", type=int, default=5)
    parser.add_argument("--start_epoch", type=int, default=0)
    parser.add_argument("--save_freq", type=int, default=10)
    parser.add_argument("--val_freq", type=int, default=1)
    parser.add_argument("--gpu", default=0, type=int)

    # Regularization
    parser.add_argument("--label_smoothing", type=float, default=0.0)
    parser.add_argument("--dropout", type=float, default=0.5)
    parser.add_argument("--mixup_alpha", type=float, default=0.0)
    parser.add_argument("--early_stopping_patience", type=int, default=0)

    # ArcFace loss
    parser.add_argument("--use_arcface", action="store_true", default=False)
    parser.add_argument("--arcface_s", type=float, default=30.0)
    parser.add_argument("--arcface_m", type=float, default=0.50)

    # Augmentation
    parser.add_argument("--aug_scale_coords", action="store_true", default=False)
    parser.add_argument("--aug_rotate", action="store_true", default=False)
    parser.add_argument("--aug_jitter_coords", action="store_true", default=False)
    parser.add_argument("--aug_shift_coords", action="store_true", default=False)
    parser.add_argument("--aug_flip", action="store_true", default=False)
    parser.add_argument("--aug_mask_feats", action="store_true", default=False)
    parser.add_argument("--aug_jitter_length", action="store_true", default=False)
    parser.add_argument("--aug_elasticate", action="store_true", default=False)
    parser.add_argument("--aug_drop_tree", action="store_true", default=False)
    parser.add_argument("--aug_skip_parent_node", action="store_true", default=False)
    parser.add_argument("--aug_swap_sibling_subtrees", action="store_true", default=False)

    # Persistence augmentation
    parser.add_argument("--use_persistence_aug", action="store_true", default=False)
    parser.add_argument("--pers_translation_scale", type=float, default=0.05)
    parser.add_argument("--pers_noise_scale", type=float, default=0.02)
    parser.add_argument("--pers_sigma_min", type=float, default=12.0)
    parser.add_argument("--pers_sigma_max", type=float, default=20.0)
    parser.add_argument("--sigma_px", type=float, default=16.0)

    # Evaluation
    parser.add_argument("--eval_mode", type=str, default="accuracy",
                        choices=["accuracy", "knn"])
    parser.add_argument("--knn_k", type=int, default=20)

    parser.add_argument("--cache_images", action="store_true", default=True)
    parser.add_argument("--debug", action="store_true", default=False)

    args = parser.parse_args()
    set_seed(args.seed)

    if args.linear_probe_epochs > 0:
        args.freeze_encoders = True
        
    # Setup work directory
    args.work_dir = f"{args.work_dir}/{args.exp_name}"
    if not os.path.exists(args.work_dir):
        os.makedirs(args.work_dir)

    # Logger
    timestamp = time.strftime("%Y%m%d_%H%M%S", time.localtime())
    if args.debug:
        log_file = None
        args.save_freq = 10000
        args.val_freq = 1
    else:
        log_file = f"{args.work_dir}/finetune_{timestamp}.log"
    logger = get_root_logger(log_file=log_file, log_level="INFO")

    logger.info("=" * 60)
    logger.info("GraPHFormer FINE-TUNING")
    logger.info(f"Mode: {args.mode}")
    if args.mode == "multimodal":
        logger.info(f"Fusion Mode: {args.fusion_mode}")
    logger.info(f"Freeze Encoders: {args.freeze_encoders}")
    logger.info(f"Pretrained Checkpoint: {args.pretrained_checkpoint}")
    logger.info(f"Dataset: {args.dataset}")
    logger.info("=" * 60)
    logger.info(json.dumps(vars(args), indent=4, sort_keys=True))

    device = torch.device("cuda")

    # Load pretrained model or create from scratch
    if args.pretrained_checkpoint is not None:
        logger.info("=> Loading pretrained model...")
        if not os.path.isfile(args.pretrained_checkpoint):
            raise FileNotFoundError(f"Checkpoint not found: {args.pretrained_checkpoint}")

        checkpoint = torch.load(args.pretrained_checkpoint, map_location=f"cuda:{args.gpu}")
        state_dict = checkpoint["state_dict"]

        # Auto-detect image encoder
        if args.mode in ['image_only', 'multimodal']:
            if "image_encoder.encoder.backbone.cls_token" in state_dict:
                args.image_encoder = "dinov2_vits14"
                logger.info(f"=> Detected DINOv2 image encoder")

        pretrained_model = CLIPModel(args).to(device)
        missing_keys, unexpected_keys = pretrained_model.load_state_dict(state_dict, strict=False)

        if missing_keys:
            logger.warning(f"=> Missing keys: {len(missing_keys)}")
        if unexpected_keys:
            logger.warning(f"=> Unexpected keys: {len(unexpected_keys)}")

        logger.info(f"=> Loaded checkpoint from epoch {checkpoint.get('epoch', 'unknown')}")
    else:
        logger.info("=> Training from scratch")
        pretrained_model = CLIPModel(args).to(device)

    # Setup augmentations
    aug_switchs = [
        False,
        args.aug_scale_coords,
        args.aug_rotate,
        args.aug_jitter_coords,
        args.aug_shift_coords,
        args.aug_flip,
        args.aug_mask_feats,
        args.aug_jitter_length,
        args.aug_elasticate,
    ]
    aug_fns = [
        None,
        RandomScaleCoords(p=0.2),
        RandomRotate(p=0.5),
        RandomJitter(p=0.2),
        RandomShift(p=0.2),
        RandomFlip(p=1),
        RandomMaskFeats(p=0.2),
        RandomJitterLength(p=0.2),
        RandomElasticate(p=0.2),
    ]
    feat_augs = [aug_fns[i] for i in range(len(aug_switchs)) if aug_switchs[i] and aug_fns[i] is not None]
    feat_augs = Compose(feat_augs) if feat_augs else None

    topo_aug_switchs = [
        args.aug_drop_tree,
        args.aug_skip_parent_node,
        args.aug_swap_sibling_subtrees,
    ]
    topo_aug_fns = [
        RandomDropSubTrees(probs=[0.05], max_cnt=5),
        RandomSkipParentNode(probs=[0.05], max_cnt=10),
        RandomSwapSiblingSubTrees(probs=[0.05], max_cnt=10),
    ]
    topo_augs = [topo_aug_fns[i] for i in range(len(topo_aug_switchs)) if topo_aug_switchs[i]]
    topo_augs = Compose(topo_augs) if topo_augs else None

    # Persistence augmentation
    persistence_aug = None
    if args.use_persistence_aug:
        persistence_aug = CombinedPersistenceAugmentation(
            translation_scale=args.pers_translation_scale,
            noise_scale=args.pers_noise_scale,
            sigma_min=args.pers_sigma_min,
            sigma_max=args.pers_sigma_max,
        )

    # Create datasets
    collate_fn = get_collate_fn(device, use_images=True)

    trainset = NeuronTreeDataset(
        phase="train",
        dataset=args.dataset,
        label_dict=LABEL_DICT[args.dataset],
        input_features=args.input_features,
        topology_transformations=topo_augs,
        attribute_transformations=feat_augs,
        use_images=True,
        image_size=args.image_size,
        cache_images=args.cache_images,
        persistence_augmentation=persistence_aug,
        sigma_px=args.sigma_px,
    )

    testset = NeuronTreeDataset(
        phase="test",
        dataset=args.dataset,
        label_dict=LABEL_DICT[args.dataset],
        input_features=args.input_features,
        use_images=True,
        image_size=args.image_size,
        cache_images=args.cache_images,
    )

    train_loader = DataLoader(
        dataset=trainset,
        batch_size=args.batch_size,
        collate_fn=collate_fn,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )

    test_loader = DataLoader(
        dataset=testset,
        batch_size=args.batch_size,
        collate_fn=collate_fn,
        shuffle=False,
        num_workers=4,
        pin_memory=True,
    )

    logger.info(f"Train samples: {len(trainset)}, Test samples: {len(testset)}")
    logger.info(f"Number of classes: {len(trainset.classes)}")

    # Create fine-tuning model
    logger.info("=> Creating fine-tuning model...")
    model = FineTuneModel(
        pretrained_model=pretrained_model,
        num_classes=len(trainset.classes),
        mode=args.mode,
        freeze_encoders=args.freeze_encoders,
        fusion_mode=args.fusion_mode if args.mode == "multimodal" else None,
        dropout=args.dropout,
        label_smoothing=args.label_smoothing,
        use_projection=args.use_projection,
        use_arcface=args.use_arcface,
        arcface_s=args.arcface_s,
        arcface_m=args.arcface_m,
        freeze_image_only=args.freeze_image_only
    ).to(device)

    del pretrained_model
    logger.info(model)

    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f"Total parameters: {total_params:,}")
    logger.info(f"Trainable parameters: {trainable_params:,}")

    # Optimizer
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=args.lr,
        weight_decay=args.wd,
    )

    cudnn.benchmark = True

    # Scheduler
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
        optimizer,
        T_0=25,
        T_mult=1,
        eta_min=args.lr * 0.01
    )

    # Training loop
    best_metric = 0.0
    best_epoch = 0
    patience_counter = 0
    total_iters = len(train_loader) * args.epochs
    current_iter = 0
    start_time = time.time()

    logger.info("=> Starting fine-tuning...")
    logger.info(f"=> Regularization: dropout={args.dropout}, label_smoothing={args.label_smoothing}")

    for epoch in range(args.start_epoch + 1, args.epochs + 1):

        if epoch == args.linear_probe_epochs + 1 and args.linear_probe_epochs > 0:
            logger.info("="*30)
            logger.info(f"Linear probe phase completed. Unfreezing encoders for full fine-tuning.")
            model.unfreeze_encoders()
            
            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            logger.info(f"Trainable parameters after unfreezing: {trainable_params:,}")

            optimizer = torch.optim.AdamW(
                filter(lambda p: p.requires_grad, model.parameters()),
                lr=args.lr / 10,
                weight_decay=args.wd,
            )
            logger.info(f"Optimizer reset to include unfrozen parameters with a new learning rate of {args.lr / 10}.")
            logger.info("="*30)

        model.train()
        epoch_loss = 0.0
        correct = 0
        total = 0

        for step, batch in enumerate(train_loader):
            try:
                batch = batch.to(device)

                if args.mixup_alpha > 0:
                    _, _, features = model(batch, return_features=True)
                    labels = batch.label.cuda() if not batch.label.is_cuda else batch.label

                    features, labels_a, labels_b, lam = mixup_data(features, labels, args.mixup_alpha)

                    if model.use_arcface:
                        extracted_features = model.feature_extractor(features)
                        logits = model.arcface(extracted_features, labels)
                        loss = model.criterion(logits, labels)
                    else:
                        logits = model.classifier(features)
                        loss = mixup_criterion(model.criterion, logits, labels_a, labels_b, lam)
                else:
                    loss, logits = model(batch)

                optimizer.zero_grad()
                loss.backward()

                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

                optimizer.step()

                pred = logits.argmax(dim=1)
                labels = batch.label.cuda() if not batch.label.is_cuda else batch.label
                correct += (pred == labels).sum().item()
                total += labels.size(0)

                epoch_loss += loss.item()
                current_iter += 1

                if step % 10 == 0:
                    current_time = time.time()
                    elapsed = current_time - start_time

                    log_str = (
                        f"Epoch {epoch:03d} | Step {step:03d}/{len(train_loader)} | "
                        f"Loss {loss.item():.4f} | "
                        f"LR {optimizer.param_groups[0]['lr']:.6f} | "
                        f"Elapsed {str(datetime.timedelta(seconds=int(elapsed)))}"
                    )
                    logger.info(log_str)
            except Exception as e:
                logger.info(f"Error in step {step}: {e}")
                continue

        scheduler.step()

        avg_loss = epoch_loss / len(train_loader)
        train_acc = correct / total * 100
        logger.info(f"Epoch {epoch:03d} | Avg Loss: {avg_loss:.4f} | Train Acc: {train_acc:.2f}%")

        # Evaluation
        if epoch % args.val_freq == 0:
            logger.info("=> Evaluating...")

            if args.eval_mode == "accuracy":
                test_acc = evaluate_accuracy(model, test_loader, device)
                logger.info(f"  Test Accuracy: {test_acc:.2f}%")
                metric = test_acc
            else:
                knn_acc = evaluate_knn(model, train_loader, test_loader, device, args.knn_k)
                logger.info(f"  KNN Accuracy (k={args.knn_k}): {knn_acc:.2f}%")
                metric = knn_acc

            if metric > best_metric:
                best_metric = metric
                best_epoch = epoch
                patience_counter = 0

                checkpoint_path = f"{args.work_dir}/best_model.pth"
                save_checkpoint(
                    {
                        "epoch": epoch,
                        "state_dict": model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "metric": metric,
                        "mode": args.mode,
                    },
                    is_best=True,
                    filename=checkpoint_path,
                )
                logger.info(f"  Saved new best checkpoint: {checkpoint_path}")
            else:
                patience_counter += 1

            logger.info(f"  Best: {best_metric:.2f}% at epoch {best_epoch}")

            if args.early_stopping_patience > 0 and patience_counter >= args.early_stopping_patience:
                logger.info(f"  Early stopping triggered")
                break

        # Save periodic checkpoint
        if epoch % args.save_freq == 0:
            checkpoint_path = f"{args.work_dir}/epoch_{epoch}.pth"
            save_checkpoint(
                {
                    "epoch": epoch,
                    "state_dict": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                },
                is_best=False,
                filename=checkpoint_path,
            )
            logger.info(f"Saved checkpoint: {checkpoint_path}")

    logger.info("Fine-tuning complete!")
    logger.info(f"Best {args.eval_mode}: {best_metric:.2f}% at epoch {best_epoch}")