src/transfer.py

import argparse
from pathlib import Path
import numpy as np
import torch

from train import (
    read_vid_mos_csv,
    split_rows,
    VQADataset,
    run_epoch,
    build_scheduler,
)
from model.qd_model import QD_MODEL


def load_pretrained_weights(model, pretrained_path, device):
    p = Path(pretrained_path)
    obj = torch.load(str(p), map_location=device, weights_only=True)
    # *.pt: dict checkpoint
    if isinstance(obj, dict) and "model" in obj:
        model.load_state_dict(obj["model"], strict=True)
        return obj
    # best_weights.pt: state_dict
    else:
        model.load_state_dict(obj, strict=True)
        return None

def make_loaders(rows_train, rows_val, rows_test, args, mos_mean, mos_std, device):
    ds_train = VQADataset(
        rows_train, args.db_path,
        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
        mos_mean=mos_mean, mos_std=mos_std,
    )
    ds_val = VQADataset(
        rows_val, args.db_path,
        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
        mos_mean=mos_mean, mos_std=mos_std,
    )
    ds_test = VQADataset(
        rows_test, args.db_path,
        clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
        mos_mean=mos_mean, mos_std=mos_std,
    )

    pin = str(device).startswith("cuda")
    loader_train = torch.utils.data.DataLoader(
        ds_train, batch_size=args.batch_size, shuffle=True,
        num_workers=args.num_workers, pin_memory=pin,
        persistent_workers=(args.num_workers > 0),
        prefetch_factor=4 if args.num_workers > 0 else None,
        drop_last=False,
    )
    loader_val = torch.utils.data.DataLoader(
        ds_val, batch_size=args.batch_size, shuffle=False,
        num_workers=args.num_workers, pin_memory=pin, drop_last=False,
    )
    loader_test = torch.utils.data.DataLoader(
        ds_test, batch_size=args.batch_size, shuffle=False,
        num_workers=args.num_workers, pin_memory=pin, drop_last=False,
    )
    return loader_train, loader_val, loader_test


def main():
    ap = argparse.ArgumentParser()
    # ----- mode -----
    ap.add_argument("--mode", choices=["finetune", "test_only"], required=True)
    ap.add_argument("--pretrained", default="/home/xinyi/Project/FD-VQA/src/checkpoints/kvq/qd_model.best.pt", help="pretrain model path")
    # ----- data -----
    ap.add_argument("--csv_path", default="/home/xinyi/Project/FD-VQA/metadata/SHORTS-SDR-DATASET_metadata.csv")
    ap.add_argument("--db_path", default="/media/xinyi/server/video_dataset/shorts-hdr-dataset/sdr/")
    ap.add_argument("--split_seed", type=int, default=0)
    ap.add_argument("--test_ratio", type=float, default=0.2)
    ap.add_argument("--val_ratio", type=float, default=0.1)  # train 80%
    # ----- video processing -----
    ap.add_argument("--clip_len", type=int, default=16)
    ap.add_argument("--resize", type=int, default=224)
    ap.add_argument("--win", type=int, default=6)
    ap.add_argument("--win_step", type=int, default=1)
    # ----- runtime -----
    ap.add_argument("--batch_size", type=int, default=8)
    ap.add_argument("--num_workers", type=int, default=4)
    ap.add_argument("--device", type=str, default="cuda")
    ap.add_argument("--no_amp", action="store_true")
    # ----- finetune hyperparams -----
    ap.add_argument("--epochs", type=int, default=10)
    ap.add_argument("--warmup_epochs", type=int, default=1)
    ap.add_argument("--lr", type=float, default=1e-5)
    ap.add_argument("--min_lr", type=float, default=1e-6)
    ap.add_argument("--finetune_lr", type=float, default=5e-5)
    ap.add_argument("--weight_decay", type=float, default=1e-2)
    ap.add_argument("--clip_unfreeze_blocks", type=int, default=4)
    ap.add_argument("--finetune_last_stage", action="store_true")
    ap.add_argument("--patience", type=int, default=6)
    # ----- save -----
    ap.add_argument("--save_dir", type=str, default="checkpoints_transfer")
    ap.add_argument("--save_name", type=str, default="transfer.pt")

    # ----- MOS normalization control -----
    # finetune: use target train mean/std
    ap.add_argument("--test_only_norm", choices=["none", "use_source_ckpt"], default="use_source_ckpt")

    args = ap.parse_args()
    torch.manual_seed(args.split_seed)
    device = torch.device(args.device)
    amp = not bool(args.no_amp)

    # ----------------------------
    # Load target rows and split
    # ----------------------------
    rows = read_vid_mos_csv(args.csv_path)
    if args.mode == "finetune":
        csv_path = Path(args.csv_path)
        if csv_path.name == "KVQ_TRAIN_metadata.csv":
            # KVQ challenge split
            val_csv = csv_path.parent / "KVQ_VAL_metadata.csv"
            test_csv = csv_path.parent / "KVQ_TEST_metadata.csv"
            train_rows = read_vid_mos_csv(str(csv_path))
            val_rows = read_vid_mos_csv(str(val_csv))
            test_rows = read_vid_mos_csv(str(test_csv))
            print(f"[KVQ split] train={len(train_rows)} val={len(val_rows)} test={len(test_rows)}")
        else:
            train_rows, val_rows, test_rows = split_rows(
                rows, seed=args.split_seed, test_ratio=args.test_ratio, val_ratio=args.val_ratio
            )
        mos_train = np.array([mos for _vid, mos in train_rows], dtype=np.float32)
        mos_mean = float(mos_train.mean()) if len(mos_train) else 0.0
        mos_std = float(mos_train.std()) if len(mos_train) else 1.0
        if mos_std <= 1e-8:
            mos_std = 1.0
    else:
        # test_only
        train_rows, val_rows, test_rows = [], [], rows
        mos_mean, mos_std = None, None

    # ----------------------------
    # Build model, load pretrained
    # ----------------------------
    model = QD_MODEL(
        clip_model="openai/clip-vit-base-patch16",
    ).to(device)

    ckpt = load_pretrained_weights(model, args.pretrained, device)
    print(f"Loaded pretrained: {args.pretrained}")

    # ----------------------------
    # Decide normalization for test_only
    # ----------------------------
    if args.mode == "test_only":
        if args.test_only_norm == "use_source_ckpt" and isinstance(ckpt, dict):
            # 注：这样对 SRCC/PLCC 没影响
            mos_mean = ckpt.get("mos_mean", None)
            mos_std = ckpt.get("mos_std", None)
            if mos_mean is None or mos_std is None:
                mos_mean, mos_std = None, None
                print("[warn] pretrained ckpt has no mos_mean/std, fallback to no normalization.")
            else:
                print(f"test_only uses source mos_mean/std from ckpt: mean={mos_mean:.4f}, std={mos_std:.4f}")
        else:
            print("test_only uses no MOS normalization.")

    # ----------------------------
    # Mode: test_only
    # ----------------------------
    if args.mode == "test_only":
        ds_test = VQADataset(
            test_rows, args.db_path,
            clip_len=args.clip_len, size=args.resize, win=args.win, win_step=args.win_step,
            mos_mean=mos_mean, mos_std=mos_std,
        )
        pin = str(device).startswith("cuda")
        loader_test = torch.utils.data.DataLoader(
            ds_test, batch_size=args.batch_size, shuffle=False,
            num_workers=args.num_workers, pin_memory=pin, drop_last=False,
        )
        print("num_test_rows =", len(test_rows))
        print("len(ds_test)  =", len(ds_test))
        print("len(loader_test) =", len(loader_test), "batch_size =", args.batch_size)

        te_loss, te_plcc, te_srcc, te_rmse = run_epoch(
            model, loader_test, device,
            optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
            desc="TestOnly", show_pbar=True
        )
        print(f"TEST_ONLY | loss={te_loss:.4f} plcc={te_plcc:.4f} srcc={te_srcc:.4f} rmse={te_rmse:.4f}")
        return

    # ----------------------------
    # DataLoaders
    # ----------------------------
    loader_train, loader_val, loader_test = make_loaders(
        train_rows, val_rows, test_rows, args,
        mos_mean=mos_mean, mos_std=mos_std, device=device
    )

    # ----------------------------
    # Mode: finetune
    # ----------------------------
    model.freeze_clip_all()
    did_unfreeze = False

    clip_params_all = []
    other_params_all = []
    for name, p in model.named_parameters():
        if name.startswith("encoder."):
            clip_params_all.append(p)
        else:
            other_params_all.append(p)

    param_groups = []
    if other_params_all:
        param_groups.append({"params": other_params_all, "lr": float(args.lr)})
    if clip_params_all:
        param_groups.append({"params": clip_params_all, "lr": float(args.finetune_lr)})

    optim = torch.optim.AdamW(param_groups, weight_decay=float(args.weight_decay))
    scheduler = build_scheduler(optim, args)

    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    last_path = save_dir / args.save_name
    best_path = save_dir / (Path(args.save_name).stem + ".best.pt")
    best_weights_path = save_dir / (Path(args.save_name).stem + ".best_weights.pt")

    best_val_srcc = -1e18
    bad_epochs = 0

    for epoch in range(1, int(args.epochs) + 1):
        if (not did_unfreeze) and bool(args.finetune_last_stage) and epoch == (int(args.warmup_epochs) + 1):
            model.unfreeze_clip_last_blocks(n_blocks=int(args.clip_unfreeze_blocks), also_unfreeze_ln=True)
            did_unfreeze = True
            print(f"[Finetune] Unfroze CLIP last {int(args.clip_unfreeze_blocks)} blocks")

        tr_loss, tr_plcc, tr_srcc, tr_rmse = run_epoch(
            model, loader_train, device,
            optim=optim, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
            desc=f"FT Train e{epoch:03d}", show_pbar=True
        )
        va_loss, va_plcc, va_srcc, va_rmse = run_epoch(
            model, loader_val, device,
            optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
            desc=f"FT Val   e{epoch:03d}", show_pbar=True
        )
        scheduler.step()

        print(
            f"epoch {epoch:03d} | "
            f"train: loss={tr_loss:.4f} plcc={tr_plcc:.4f} srcc={tr_srcc:.4f} rmse={tr_rmse:.4f} | "
            f"val: loss={va_loss:.4f} plcc={va_plcc:.4f} srcc={va_srcc:.4f} rmse={va_rmse:.4f}"
        )

        ckpt_out = {
            "epoch": epoch,
            "model": model.state_dict(),
            "optim": optim.state_dict(),
            "mos_mean": mos_mean,
            "mos_std": mos_std,
            "args": vars(args),
            "best_val_srcc": best_val_srcc,
        }
        torch.save(ckpt_out, str(last_path))

        if va_srcc > best_val_srcc:
            best_val_srcc = va_srcc
            bad_epochs = 0
            ckpt_out["best_val_srcc"] = best_val_srcc
            torch.save(ckpt_out, str(best_path))
            torch.save(model.state_dict(), str(best_weights_path))
            print(f"  [best] val_srcc={best_val_srcc:.4f} -> saved {best_weights_path}")
        else:
            bad_epochs += 1
            if bad_epochs >= int(args.patience):
                print(f"[EarlyStop] val_srcc not improved for {bad_epochs} epochs. Stop.")
                break

    # load best and test
    if best_weights_path.exists():
        sd = torch.load(str(best_weights_path), map_location=device, weights_only=True)
        model.load_state_dict(sd, strict=True)
        print(f"Loaded best weights: {best_weights_path}")

    te_loss, te_plcc, te_srcc, te_rmse = run_epoch(
        model, loader_test, device,
        optim=None, amp=amp, mos_mean=mos_mean, mos_std=mos_std,
        desc="FT Test", show_pbar=True
    )
    print(f"FINETUNE TEST | loss={te_loss:.4f} plcc={te_plcc:.4f} srcc={te_srcc:.4f} rmse={te_rmse:.4f}")


if __name__ == "__main__":
    main()