train.py

from __future__ import annotations

import argparse
import json
import math
import sys
from pathlib import Path

import open_clip
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT / "src"))

from data import SUPPORTED_DATASETS, ReIDDataset, collate_reid, load_reid_records
from training import (
    AverageMeter,
    build_optimizer,
    build_warmup_cosine_scheduler,
    contrastive_step,
    evaluate,
    get_device,
    move_batch_to_device,
    seed_everything,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Full fine-tune OpenCLIP on text-based person ReID datasets with contrastive learning."
    )
    parser.add_argument(
        "--env-file",
        default="env/.env",
        help="Path to the .env file containing DATASET_ROOT.",
    )
    parser.add_argument("--dataset", required=True, choices=SUPPORTED_DATASETS)
    parser.add_argument("--model-name", default="ViT-B-16", help="OpenCLIP model name.")
    parser.add_argument(
        "--pretrained",
        default="laion2b_s34b_b88k",
        help="OpenCLIP pretrained weight tag, for example openai or laion2b_s34b_b88k.",
    )
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--weight-decay", type=float, default=0.2)
    parser.add_argument("--warmup-ratio", type=float, default=0.05)
    parser.add_argument("--grad-clip-norm", type=float, default=1.0)
    parser.add_argument("--accum-steps", type=int, default=1)
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"])
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--caption-mode", default="all", choices=["all", "random"])
    parser.add_argument("--train-split", default="train")
    parser.add_argument("--val-split", default="val")
    parser.add_argument(
        "--max-train-samples", type=int, default=0, help="Debug limit. 0 means use all."
    )
    parser.add_argument(
        "--max-val-samples", type=int, default=0, help="Debug limit. 0 means use all."
    )
    parser.add_argument(
        "--eval-every", type=int, default=1, help="Run validation every N epochs."
    )
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--resume", default="", help="Checkpoint path to resume from.")
    parser.add_argument(
        "--no-amp", action="store_true", help="Disable CUDA mixed precision."
    )
    parser.add_argument(
        "--freeze-backbone",
        action="store_true",
        help="Freeze all parameters except visual.proj and text_projection.",
    )
    parser.add_argument(
        "--no-grad-checkpointing",
        action="store_true",
        help="Disable gradient checkpointing (enabled by default).",
    )
    parser.add_argument(
        "--save-every",
        type=int,
        default=0,
        help="Also save every N epochs. 0 disables.",
    )
    return parser.parse_args()


def make_loader(
    dataset: ReIDDataset,
    batch_size: int,
    num_workers: int,
    tokenizer,
    shuffle: bool,
    drop_last: bool,
) -> DataLoader:
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=torch.cuda.is_available(),
        drop_last=drop_last,
        persistent_workers=num_workers > 0,
        collate_fn=lambda batch: collate_reid(batch, tokenizer),
    )


def save_checkpoint(
    path: Path,
    *,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    scaler: torch.amp.GradScaler,
    args: argparse.Namespace,
    epoch: int,
    global_step: int,
    best_val_score: float,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "scaler_state_dict": scaler.state_dict(),
            "args": vars(args),
            "epoch": epoch,
            "global_step": global_step,
            "best_val_score": best_val_score,
        },
        path,
    )


def load_checkpoint(
    path: Path,
    *,
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    scaler: torch.amp.GradScaler,
    device: torch.device,
) -> tuple[int, int, float]:
    checkpoint = torch.load(path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
    if checkpoint.get("scaler_state_dict"):
        scaler.load_state_dict(checkpoint["scaler_state_dict"])
    start_epoch = int(checkpoint.get("epoch", 0)) + 1
    global_step = int(checkpoint.get("global_step", 0))
    best_val_score = float(checkpoint.get("best_val_score", -math.inf))
    return start_epoch, global_step, best_val_score


def main() -> None:
    args = parse_args()
    if args.accum_steps < 1:
        raise ValueError("--accum-steps must be >= 1.")
    if args.epochs < 1:
        raise ValueError("--epochs must be >= 1.")
    if args.eval_every < 1:
        raise ValueError("--eval-every must be >= 1.")

    seed_everything(args.seed)
    device = get_device(args.device)
    output_dir = Path(args.output_dir or f"artifacts/{args.dataset}_contrastive")
    output_dir.mkdir(parents=True, exist_ok=True)

    records = load_reid_records(env_file=args.env_file, dataset=args.dataset)
    train_records = [record for record in records if record.split == args.train_split]
    val_records = [record for record in records if record.split == args.val_split]
    if not train_records:
        raise RuntimeError(f"No records found for train split {args.train_split!r}.")
    if not val_records:
        raise RuntimeError(f"No records found for validation split {args.val_split!r}.")

    model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(
        args.model_name,
        pretrained=args.pretrained,
        device=device,
    )
    tokenizer = open_clip.get_tokenizer(args.model_name)
    model.train()
    if args.freeze_backbone:
        _PROJECTION_PARAMS = {"visual.proj", "text_projection"}
        for name, parameter in model.named_parameters():
            parameter.requires_grad_(name in _PROJECTION_PARAMS)
    else:
        for parameter in model.parameters():
            parameter.requires_grad_(True)
        if not args.no_grad_checkpointing:
            model.set_grad_checkpointing(True)

    train_dataset = ReIDDataset(
        train_records,
        transform=preprocess_train,
        caption_mode=args.caption_mode,
        max_samples=args.max_train_samples or None,
    )
    val_dataset = ReIDDataset(
        val_records,
        transform=preprocess_val,
        caption_mode="all",
        max_samples=args.max_val_samples or None,
    )
    train_loader = make_loader(
        train_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        tokenizer=tokenizer,
        shuffle=True,
        drop_last=True,
    )
    val_loader = make_loader(
        val_dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        tokenizer=tokenizer,
        shuffle=False,
        drop_last=False,
    )
    if len(train_loader) == 0:
        raise RuntimeError(
            "Training loader is empty. Reduce --batch-size or increase --max-train-samples because training uses drop_last."
        )

    optimizer = build_optimizer(model, lr=args.lr, weight_decay=args.weight_decay)
    total_update_steps = math.ceil(len(train_loader) / args.accum_steps) * args.epochs
    warmup_steps = int(total_update_steps * args.warmup_ratio)
    scheduler = build_warmup_cosine_scheduler(
        optimizer, warmup_steps, total_update_steps
    )
    use_amp = device.type == "cuda" and not args.no_amp
    scaler = torch.amp.GradScaler(device.type, enabled=use_amp)

    start_epoch = 0
    global_step = 0
    best_val_score = -math.inf
    if args.resume:
        start_epoch, global_step, best_val_score = load_checkpoint(
            Path(args.resume),
            model=model,
            optimizer=optimizer,
            scheduler=scheduler,
            scaler=scaler,
            device=device,
        )

    run_config_path = output_dir / "run_config.json"
    run_config_path.write_text(
        json.dumps(vars(args), indent=2, ensure_ascii=False), encoding="utf-8"
    )

    print(
        f"Dataset: {args.dataset} train={len(train_dataset):,} examples, val={len(val_dataset):,} examples"
    )
    mode = "projection-only" if args.freeze_backbone else "full fine-tuning"
    print(f"Model: {args.model_name} ({args.pretrained}) on {device}; {mode} enabled")
    print(
        f"Optimizer updates: {total_update_steps:,}, warmup={warmup_steps:,}, amp={use_amp}"
    )

    for epoch in range(start_epoch, args.epochs):
        model.train()
        loss_meter = AverageMeter()
        i2t_meter = AverageMeter()
        t2i_meter = AverageMeter()
        optimizer.zero_grad(set_to_none=True)

        progress = tqdm(
            train_loader, desc=f"epoch {epoch + 1}/{args.epochs}", dynamic_ncols=True
        )
        for step, batch in enumerate(progress, start=1):
            batch = move_batch_to_device(batch, device)
            with torch.amp.autocast(device_type=device.type, enabled=use_amp):
                output = contrastive_step(model, batch)
                loss = output.loss / args.accum_steps

            scaler.scale(loss).backward()

            should_update = step % args.accum_steps == 0 or step == len(train_loader)
            if should_update:
                if args.grad_clip_norm > 0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(
                        model.parameters(), args.grad_clip_norm
                    )
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()
                optimizer.zero_grad(set_to_none=True)
                global_step += 1

            batch_size = int(batch["images"].shape[0])
            loss_meter.update(float(output.loss.detach().cpu()), batch_size)
            i2t_meter.update(
                float(output.image_to_text_top1.detach().cpu()), batch_size
            )
            t2i_meter.update(
                float(output.text_to_image_top1.detach().cpu()), batch_size
            )
            progress.set_postfix(
                loss=f"{loss_meter.avg:.4f}",
                i2t=f"{i2t_meter.avg:.3f}",
                t2i=f"{t2i_meter.avg:.3f}",
                lr=f"{scheduler.get_last_lr()[0]:.2e}",
            )

        val_metrics = None
        if (epoch + 1) % args.eval_every == 0:
            val_metrics = evaluate(model, val_loader, device=device, use_amp=use_amp)
            val_score = (
                val_metrics["image_to_text_top1"] + val_metrics["text_to_image_top1"]
            ) / 2.0
            print(
                "validation "
                f"loss={val_metrics['loss']:.4f} "
                f"i2t@1={val_metrics['image_to_text_top1']:.4f} "
                f"t2i@1={val_metrics['text_to_image_top1']:.4f}"
            )
            if val_score > best_val_score:
                best_val_score = val_score
                save_checkpoint(
                    output_dir / "best.pt",
                    model=model,
                    optimizer=optimizer,
                    scheduler=scheduler,
                    scaler=scaler,
                    args=args,
                    epoch=epoch,
                    global_step=global_step,
                    best_val_score=best_val_score,
                )

        save_checkpoint(
            output_dir / "last.pt",
            model=model,
            optimizer=optimizer,
            scheduler=scheduler,
            scaler=scaler,
            args=args,
            epoch=epoch,
            global_step=global_step,
            best_val_score=best_val_score,
        )
        if args.save_every and (epoch + 1) % args.save_every == 0:
            save_checkpoint(
                output_dir / f"epoch_{epoch + 1}.pt",
                model=model,
                optimizer=optimizer,
                scheduler=scheduler,
                scaler=scaler,
                args=args,
                epoch=epoch,
                global_step=global_step,
                best_val_score=best_val_score,
            )

        if val_metrics is None:
            print(f"epoch {epoch + 1} train loss={loss_meter.avg:.4f}")


if __name__ == "__main__":
    main()