train_vit.py

#!/usr/bin/env python3
"""
Fine-tune a pretrained ViT-S/16 on a MedMNIST 2D dataset.

Requirements:
  pip install timm torchvision torchaudio torch medmnist

Example usages:
  # Head-only tuning (fast, low risk of overfit)
  python finetune_vit_medmnist.py --dataset pathmnist --epochs 20 --head-only

  # Last-4-blocks tuning + layer-wise LR decay
  python finetune_vit_medmnist.py --dataset pathmnist --epochs 50 --unfreeze-last 4 --lrd 0.75

  # Full fine-tune with small LR
  python finetune_vit_medmnist.py --dataset organamnist --epochs 100 --lr 5e-5 --weight-decay 0.05

Outputs:
  - best.pth  : state dict of the best (val accuracy) model
  - last.pth  : last epoch weights
"""

import argparse
import json
import logging
import math
import os
import pickle
import random
import ssl
import time
import urllib3
from datetime import datetime
from pathlib import Path
from typing import List, Tuple

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader

import timm
from timm.layers import set_fast_norm
from tqdm import tqdm

import torchvision.transforms as T

try:
    import medmnist
    from medmnist import INFO
    from medmnist.dataset import (
        PathMNIST,
        ChestMNIST,
        DermaMNIST,
        OCTMNIST,
        PneumoniaMNIST,
        RetinaMNIST,
        BreastMNIST,
        BloodMNIST,
        TissueMNIST,
        OrganAMNIST,
        OrganCMNIST,
        OrganSMNIST,
    )
except Exception as e:
    raise RuntimeError("Please 'pip install medmnist' to use this script.") from e


NAME2CLS = {
    "pathmnist": PathMNIST,
    "chestmnist": ChestMNIST,
    "dermamnist": DermaMNIST,
    "octmnist": OCTMNIST,
    "pneumoniamnist": PneumoniaMNIST,
    "retinamnist": RetinaMNIST,
    "breastmnist": BreastMNIST,
    "bloodmnist": BloodMNIST,
    "tissuemnist": TissueMNIST,
    "organamnist": OrganAMNIST,
    "organcmnist": OrganCMNIST,
    "organsmnist": OrganSMNIST,
}


def set_seed(seed: int):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True


def fix_ssl_issues():
    """Fix SSL certificate verification issues for Hugging Face downloads."""
    # Disable SSL warnings
    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

    # Create unverified SSL context
    ssl_context = ssl.create_default_context()
    ssl_context.check_hostname = False
    ssl_context.verify_mode = ssl.CERT_NONE

    # Set as default SSL context
    ssl._create_default_https_context = ssl._create_unverified_context


def setup_logging_and_outputs(args):
    """Set up logging and create output directories."""
    # Create timestamp for this run
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # Create output directory
    output_dir = Path(f"outputs/{args.dataset}_{timestamp}")
    output_dir.mkdir(parents=True, exist_ok=True)

    # Create subdirectories
    (output_dir / "checkpoints").mkdir(exist_ok=True)
    (output_dir / "logs").mkdir(exist_ok=True)

    # Set up logging
    log_file = output_dir / "logs" / "training.log"

    # Configure logging
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(),  # Also log to console
        ],
    )

    logger = logging.getLogger(__name__)
    logger.info(f"Starting training run: {timestamp}")
    logger.info(f"Output directory: {output_dir}")
    logger.info(f"Arguments: {vars(args)}")

    return output_dir, logger


def save_checkpoint(
    model,
    optimizer,
    scheduler,
    epoch,
    val_acc,
    train_loss,
    val_loss,
    output_dir,
    is_best=False,
    logger=None,
):
    """Save model checkpoint with comprehensive metadata."""
    checkpoint = {
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "scheduler_state_dict": scheduler.state_dict(),
        "val_acc": val_acc,
        "train_loss": train_loss,
        "val_loss": val_loss,
        "timestamp": datetime.now().isoformat(),
        "model_config": {
            "num_classes": model.num_classes if hasattr(model, "num_classes") else None,
            "model_name": getattr(model, "default_cfg", {}).get(
                "model_name", "unknown"
            ),
        },
    }

    # Save last checkpoint
    last_path = output_dir / "checkpoints" / "last.pth"
    torch.save(checkpoint, last_path)

    # Save best checkpoint if applicable
    if is_best:
        best_path = output_dir / "checkpoints" / "best.pth"
        torch.save(checkpoint, best_path)
        if logger:
            logger.info(f"New best model saved with val_acc={val_acc*100:.2f}%")

    if logger:
        logger.info(f"Checkpoint saved: epoch={epoch}, val_acc={val_acc*100:.2f}%")


def log_training_metrics(
    epoch, train_loss, val_loss, val_acc, lr, logger, training_history=None
):
    """Log training metrics and update history."""
    if training_history is None:
        training_history = {
            "epoch": [],
            "train_loss": [],
            "val_loss": [],
            "val_acc": [],
            "lr": [],
        }

    # Update history
    training_history["epoch"].append(epoch)
    training_history["train_loss"].append(train_loss)
    training_history["val_loss"].append(val_loss)
    training_history["val_acc"].append(val_acc)
    training_history["lr"].append(lr)

    # Log metrics
    logger.info(
        f"Epoch {epoch:03d} | train_loss={train_loss:.4f} | "
        f"val_loss={val_loss:.4f} | val_acc={val_acc*100:.2f}% | lr={lr:.2e}"
    )

    return training_history


def build_transforms(img_size=224):
    # MedMNIST images are 28x28 (grayscale); we upscale and apply moderate aug.
    train_tfms = T.Compose(
        [
            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.BICUBIC),
            T.RandomHorizontalFlip(p=0.5),
            T.RandomRotation(10),
            T.RandomResizedCrop(img_size, scale=(0.85, 1.0)),
            T.ToTensor(),
            # If grayscale, repeat to 3 channels inside dataset __getitem__.
            T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
        ]
    )
    val_tfms = T.Compose(
        [
            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.BICUBIC),
            T.CenterCrop(img_size),
            T.ToTensor(),
            T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
        ]
    )
    return train_tfms, val_tfms


class CachedRGBWrapper(torch.utils.data.Dataset):
    """Cached version of RGBWrapper that pre-processes and saves data to disk."""

    def __init__(self, base_ds, transform=None, cache_dir=None, force_rebuild=False):
        self.base = base_ds
        self.transform = transform
        self.is_multilabel = bool(
            getattr(self.base, "task", "") == "multi-label, binary-class"
        )

        # Set up caching
        if cache_dir is not None:
            self.cache_dir = Path(cache_dir)
            self.cache_dir.mkdir(parents=True, exist_ok=True)
            self.cache_file = (
                self.cache_dir / f"{self.base.__class__.__name__}_{self.base.split}.pkl"
            )

            if self.cache_file.exists() and not force_rebuild:
                print(f"Loading cached dataset from {self.cache_file}")
                self._load_cache()
            else:
                print(f"Building cache for dataset: {self.cache_file}")
                self._build_cache()
        else:
            self.cached_data = None

    def _build_cache(self):
        """Build cache by processing all data."""
        self.cached_data = []
        print(f"Processing {len(self.base)} samples...")

        # Use tqdm for progress bar
        for i in tqdm(range(len(self.base)), desc="Building cache", unit="samples"):
            item = self.base[i]
            img, label = item[0], item[1]

            # Process image
            if isinstance(img, torch.Tensor):
                x = img
            elif hasattr(img, "mode"):  # PIL Image
                x = T.ToTensor()(img)
            else:
                x = torch.tensor(img, dtype=torch.float32)

            # Ensure proper shape: CxHxW
            if x.ndim == 2:
                x = x.unsqueeze(0)  # 1xHxW
            elif x.ndim == 3 and x.shape[0] > 3:
                x = x.permute(2, 0, 1)

            # Convert to 3 channels if needed
            if x.shape[0] == 1:
                x = x.repeat(3, 1, 1)  # 3xHxW
            elif x.shape[0] == 4:  # RGBA
                x = x[:3]  # Take only RGB channels

            # Convert back to PIL for transforms
            x = T.ToPILImage()(x)
            if self.transform is not None:
                x = self.transform(x)

            # Process label
            if self.is_multilabel:
                y = torch.tensor(label, dtype=torch.float32)
                if y.dim() == 0:
                    y = y.unsqueeze(0)
            else:
                y = torch.tensor(label, dtype=torch.long)
                if y.dim() > 0:
                    y = y.flatten()[0]

            self.cached_data.append((x, y))

        # Save cache
        if hasattr(self, "cache_file"):
            with open(self.cache_file, "wb") as f:
                pickle.dump(self.cached_data, f)
            print(f"Cache saved to {self.cache_file}")

    def _load_cache(self):
        """Load data from cache."""
        with open(self.cache_file, "rb") as f:
            self.cached_data = pickle.load(f)
        print(f"Loaded {len(self.cached_data)} samples from cache")

    def __len__(self):
        return len(self.base)

    def __getitem__(self, idx):
        if self.cached_data is not None:
            return self.cached_data[idx]
        else:
            # Fallback to original processing
            return self._process_item(idx)

    def _process_item(self, idx):
        """Original item processing logic."""
        item = self.base[idx]
        img, label = item[0], item[1]

        # Handle different image types
        if isinstance(img, torch.Tensor):
            x = img
        elif hasattr(img, "mode"):  # PIL Image
            x = T.ToTensor()(img)
        else:
            x = torch.tensor(img, dtype=torch.float32)

        # Ensure proper shape: CxHxW
        if x.ndim == 2:
            x = x.unsqueeze(0)  # 1xHxW
        elif x.ndim == 3 and x.shape[0] > 3:
            x = x.permute(2, 0, 1)

        # Convert to 3 channels if needed
        if x.shape[0] == 1:
            x = x.repeat(3, 1, 1)  # 3xHxW
        elif x.shape[0] == 4:  # RGBA
            x = x[:3]  # Take only RGB channels

        # Convert back to PIL for transforms
        x = T.ToPILImage()(x)
        if self.transform is not None:
            x = self.transform(x)

        # Handle label conversion properly
        if self.is_multilabel:
            y = torch.tensor(label, dtype=torch.float32)
            if y.dim() == 0:
                y = y.unsqueeze(0)
        else:
            y = torch.tensor(label, dtype=torch.long)
            if y.dim() > 0:
                y = y.flatten()[0]
        return x, y


class ClassFilteredWrapper(torch.utils.data.Dataset):
    """Wrapper that filters out a specific class and remaps remaining labels to be consecutive.

    Example: If removing class 0 from a dataset with classes [0,1,2,3,4]:
    - Filters out all samples with label 0
    - Remaps: 1->0, 2->1, 3->2, 4->3
    """

    def __init__(self, base_ds, remove_class=0):
        self.base = base_ds
        self.remove_class = remove_class

        # Find all valid indices (excluding samples with remove_class)
        self.valid_indices = []
        for i in range(len(self.base)):
            item = self.base[i]
            label = item[1]
            # Handle different label formats
            if isinstance(label, torch.Tensor):
                label_val = (
                    label.item() if label.numel() == 1 else label.flatten()[0].item()
                )
            elif isinstance(label, (list, tuple)) and len(label) > 0:
                label_val = label[0]
            else:
                label_val = int(label)

            if label_val != remove_class:
                self.valid_indices.append(i)

        # Create label mapping: old_label -> new_label
        # Get unique labels from valid samples
        valid_labels = set()
        for idx in self.valid_indices:
            item = self.base[idx]
            label = item[1]
            if isinstance(label, torch.Tensor):
                label_val = (
                    label.item() if label.numel() == 1 else label.flatten()[0].item()
                )
            elif isinstance(label, (list, tuple)) and len(label) > 0:
                label_val = label[0]
            else:
                label_val = int(label)
            valid_labels.add(label_val)

        # Create mapping: sort labels and map to 0, 1, 2, ...
        sorted_labels = sorted(valid_labels)
        self.label_mapping = {
            old_label: new_label for new_label, old_label in enumerate(sorted_labels)
        }

        print(f"ClassFilteredWrapper: Removed class {remove_class}")
        print(f"  Original dataset size: {len(self.base)}")
        print(f"  Filtered dataset size: {len(self.valid_indices)}")
        print(f"  Label mapping: {self.label_mapping}")

    def __len__(self):
        return len(self.valid_indices)

    def __getitem__(self, idx):
        # Get the actual index in the base dataset
        actual_idx = self.valid_indices[idx]
        item = self.base[actual_idx]
        img, label = item[0], item[1]

        # Get original label value
        if isinstance(label, torch.Tensor):
            label_val = (
                label.item() if label.numel() == 1 else label.flatten()[0].item()
            )
        elif isinstance(label, (list, tuple)) and len(label) > 0:
            label_val = label[0]
        else:
            label_val = int(label)

        # Remap label
        new_label = self.label_mapping[label_val]
        new_label = torch.tensor(new_label, dtype=torch.long)

        return img, new_label


class RGBWrapper(torch.utils.data.Dataset):
    """Wrap a MedMNIST dataset to ensure 3-channel output and (img, label) tuples.
    MedMNIST returns dict-like; we convert and repeat channels if needed.
    """

    def __init__(self, base_ds, transform=None):
        self.base = base_ds
        self.transform = transform
        self.is_multilabel = bool(
            getattr(self.base, "task", "") == "multi-label, binary-class"
        )

    def __len__(self):
        return len(self.base)

    def __getitem__(self, idx):
        item = self.base[idx]
        img, label = item[0], item[1]

        # Handle different image types
        if isinstance(img, torch.Tensor):
            x = img
        elif hasattr(img, "mode"):  # PIL Image
            # Convert PIL Image to tensor
            x = T.ToTensor()(img)
        else:
            # numpy array -> tensor
            x = torch.tensor(img, dtype=torch.float32)

        # Ensure proper shape: CxHxW
        if x.ndim == 2:
            x = x.unsqueeze(0)  # 1xHxW
        elif x.ndim == 3 and x.shape[0] > 3:
            # If channels are last (HxWxC), transpose to (CxHxW)
            x = x.permute(2, 0, 1)

        # Convert to 3 channels if needed
        if x.shape[0] == 1:
            x = x.repeat(3, 1, 1)  # 3xHxW
        elif x.shape[0] == 4:  # RGBA
            x = x[:3]  # Take only RGB channels

        # Convert back to PIL for transforms
        x = T.ToPILImage()(x)
        if self.transform is not None:
            x = self.transform(x)

        # Handle label conversion properly
        if self.is_multilabel:
            y = torch.tensor(label, dtype=torch.float32)
            if y.dim() == 0:  # scalar
                y = y.unsqueeze(0)
        else:
            y = torch.tensor(label, dtype=torch.long)
            if y.dim() > 0:  # multi-dimensional, take the first element
                y = y.flatten()[0]
        return x, y


def get_dataloaders(
    name: str,
    img_size: int,
    batch_size: int,
    workers: int,
    cache_dir: str = None,
    use_cache: bool = True,
    force_rebuild: bool = False,
    remove_class: int = None,
):
    name = name.lower()
    assert (
        name in NAME2CLS
    ), f"Unknown dataset '{name}'. Available: {list(NAME2CLS.keys())}"
    info = INFO[name]
    n_classes = (
        len(info["label"]) if isinstance(info["label"], dict) else info["n_classes"]
    )
    train_tfms, val_tfms = build_transforms(img_size)

    DS = NAME2CLS[name]
    train_ds = DS(split="train", download=True)
    val_ds = DS(split="val", download=True)
    test_ds = DS(split="test", download=True)

    # Use cached wrapper if enabled
    if use_cache and cache_dir is not None:
        train_ds = CachedRGBWrapper(
            train_ds,
            transform=train_tfms,
            cache_dir=cache_dir,
            force_rebuild=force_rebuild,
        )
        val_ds = CachedRGBWrapper(
            val_ds, transform=val_tfms, cache_dir=cache_dir, force_rebuild=force_rebuild
        )
        test_ds = CachedRGBWrapper(
            test_ds,
            transform=val_tfms,
            cache_dir=cache_dir,
            force_rebuild=force_rebuild,
        )
    else:
        train_ds = RGBWrapper(train_ds, transform=train_tfms)
        val_ds = RGBWrapper(val_ds, transform=val_tfms)
        test_ds = RGBWrapper(test_ds, transform=val_tfms)

    # Filter out specified class if requested
    if remove_class is not None:
        train_ds = ClassFilteredWrapper(train_ds, remove_class=remove_class)
        val_ds = ClassFilteredWrapper(val_ds, remove_class=remove_class)
        test_ds = ClassFilteredWrapper(test_ds, remove_class=remove_class)
        # Update number of classes
        n_classes = n_classes - 1
        print(f"After removing class {remove_class}, updated n_classes = {n_classes}")

    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=workers,
        pin_memory=True,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=workers,
        pin_memory=True,
    )
    test_loader = DataLoader(
        test_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=workers,
        pin_memory=True,
    )
    is_multilabel = bool(INFO[name]["task"] == "multi-label, binary-class")
    return train_loader, val_loader, test_loader, n_classes, is_multilabel


def make_model(
    n_classes: int, drop_path: float = 0.1, dropout: float = 0.1, offline: bool = False
):
    """Create a Vision Transformer model with SSL issue handling."""

    # Fix SSL issues before attempting download
    fix_ssl_issues()

    # Try different approaches to load the model
    model = None
    attempts = [
        # Attempt 1: Try with pretrained weights
        lambda: timm.create_model(
            "vit_small_patch16_224",
            pretrained=True,
            num_classes=n_classes,
            drop_path_rate=drop_path,
            drop_rate=dropout,
        ),
        # Attempt 2: Try without pretrained weights if offline or download fails
        lambda: timm.create_model(
            "vit_small_patch16_224",
            pretrained=False,
            num_classes=n_classes,
            drop_path_rate=drop_path,
            drop_rate=dropout,
        ),
        # Attempt 3: Try with a different model variant
        lambda: timm.create_model(
            "vit_tiny_patch16_224",
            pretrained=False,
            num_classes=n_classes,
            drop_path_rate=drop_path,
            drop_rate=dropout,
        ),
    ]

    for i, attempt in enumerate(attempts):
        try:
            print(f"Attempting to create model (attempt {i+1}/{len(attempts)})...")
            model = attempt()
            print(f"Successfully created model on attempt {i+1}")
            break
        except Exception as e:
            print(f"Attempt {i+1} failed: {str(e)}")
            if i == len(attempts) - 1:
                print("All attempts failed. Creating a basic ViT model...")
                # Last resort: create a basic model
                model = timm.create_model(
                    "vit_tiny_patch16_224",
                    pretrained=False,
                    num_classes=n_classes,
                    drop_path_rate=drop_path,
                    drop_rate=dropout,
                )
                break
            continue

    return model


def param_groups_lrd(
    model: nn.Module, base_lr: float, weight_decay: float, lrd: float = 0.8
):
    """Layer-wise LR decay for ViT blocks: deeper layers get higher LR.
    lrd in (0,1]; e.g., 0.75 means each layer earlier gets LR *= 0.75.
    """
    layers: List[nn.Module] = []
    if hasattr(model, "patch_embed"):
        layers.append(model.patch_embed)
    if hasattr(model, "pos_drop"):
        layers.append(model.pos_drop)
    if hasattr(model, "blocks"):
        for blk in model.blocks:
            layers.append(blk)
    if hasattr(model, "norm"):
        layers.append(model.norm)
    # head is separate, will be added with base_lr

    total = len(layers)
    groups = []
    for i, layer in enumerate(layers):
        decay = weight_decay
        lr = base_lr * (lrd ** (total - i - 1))
        params = [p for p in layer.parameters() if p.requires_grad]
        if params:
            groups.append({"params": params, "lr": lr, "weight_decay": decay})
    # classifier head
    head_params = [p for p in model.get_classifier().parameters() if p.requires_grad]
    if head_params:
        groups.append(
            {"params": head_params, "lr": base_lr, "weight_decay": weight_decay}
        )
    return groups


def evaluate(model, loader, device, is_multilabel=False, desc="Evaluating"):
    model.eval()
    correct = 0
    total = 0
    loss_sum = 0.0
    criterion = nn.BCEWithLogitsLoss() if is_multilabel else nn.CrossEntropyLoss()
    with torch.no_grad():
        # Add progress bar for evaluation
        eval_pbar = tqdm(loader, desc=desc, leave=False, unit="batch")
        for x, y in eval_pbar:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            logits = model(x)
            loss = criterion(logits, y)
            loss_sum += loss.item() * x.size(0)
            if is_multilabel:
                preds = (torch.sigmoid(logits) > 0.5).int()
                correct += (preds == y.int()).all(dim=1).sum().item()
            else:
                pred = logits.argmax(dim=1)
                correct += (pred == y).sum().item()
            total += x.size(0)

            # Update progress bar with current accuracy
            current_acc = correct / total if total > 0 else 0
            eval_pbar.set_postfix({"acc": f"{current_acc:.3f}"})

    return loss_sum / total, correct / total


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset", type=str, required=True, choices=list(NAME2CLS.keys())
    )
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--img-size", type=int, default=224)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--lr", type=float, default=5e-4)
    parser.add_argument("--weight-decay", type=float, default=0.05)
    parser.add_argument(
        "--lrd",
        type=float,
        default=1.0,
        help="layer-wise LR decay factor; 1.0 disables LLRD",
    )
    parser.add_argument("--drop-path", type=float, default=0.1)

    parser.add_argument(
        "--head-only", action="store_true", help="train only the classifier head"
    )
    parser.add_argument(
        "--unfreeze-last",
        type=int,
        default=0,
        help="also unfreeze last N transformer blocks",
    )

    parser.add_argument(
        "--label-smoothing",
        type=float,
        default=0.1,
        help="Label smoothing factor (0.0 = no smoothing, 0.1 = moderate, 0.2+ = strong)",
    )
    parser.add_argument(
        "--mixup",
        type=float,
        default=0.2,
        help="Mixup alpha parameter (0.0 = no mixup, 0.2 = moderate, 0.4+ = strong)",
    )
    parser.add_argument("--cutmix", type=float, default=0.0)
    parser.add_argument(
        "--dropout",
        type=float,
        default=0.1,
        help="Dropout rate for additional regularization",
    )

    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--amp", action="store_true")
    parser.add_argument(
        "--offline",
        action="store_true",
        help="Use offline mode to avoid downloading pretrained weights",
    )
    parser.add_argument(
        "--cache-dir",
        type=str,
        default="data_cache",
        help="Directory to cache processed datasets",
    )
    parser.add_argument(
        "--no-cache",
        action="store_true",
        help="Disable dataset caching",
    )
    parser.add_argument(
        "--rebuild-cache",
        action="store_true",
        help="Force rebuild of dataset cache",
    )
    parser.add_argument(
        "--remove-class",
        type=int,
        default=None,
        help="Remove a specific class from the dataset (e.g., 0 to remove class 0). Remaining classes will be remapped to consecutive labels.",
    )
    args = parser.parse_args()

    # Set up logging and output directories
    output_dir, logger = setup_logging_and_outputs(args)

    set_seed(args.seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info(f"Using device: {device}")
    set_fast_norm()  # small speed win

    train_loader, val_loader, test_loader, n_classes, is_multilabel = get_dataloaders(
        args.dataset,
        args.img_size,
        args.batch_size,
        args.workers,
        cache_dir=args.cache_dir,
        use_cache=not args.no_cache,
        force_rebuild=args.rebuild_cache,
        remove_class=args.remove_class,
    )
    logger.info(
        f"Dataset: {args.dataset}, Classes: {n_classes}, Multilabel: {is_multilabel}"
    )

    model = make_model(
        n_classes, drop_path=args.drop_path, dropout=args.dropout, offline=args.offline
    )
    logger.info(
        f"Model created with {sum(p.numel() for p in model.parameters())} parameters"
    )
    logger.info(
        f"Regularization settings: label_smoothing={args.label_smoothing}, "
        f"mixup={args.mixup}, dropout={args.dropout}, drop_path={args.drop_path}"
    )

    # Freeze per config
    if args.head_only:
        for p in model.parameters():
            p.requires_grad = False
        for p in model.get_classifier().parameters():
            p.requires_grad = True
    elif args.unfreeze_last > 0:
        for p in model.parameters():
            p.requires_grad = False
        # unfreeze classifier
        for p in model.get_classifier().parameters():
            p.requires_grad = True
        # unfreeze last N transformer blocks
        N = args.unfreeze_last
        for blk in model.blocks[-N:]:
            for p in blk.parameters():
                p.requires_grad = True
        # also the final norm
        for p in model.norm.parameters():
            p.requires_grad = True
    else:
        # full fine-tune
        for p in model.parameters():
            p.requires_grad = True

    model.to(device)

    # Loss
    if is_multilabel:
        criterion = nn.BCEWithLogitsLoss()
    else:
        criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)

    # Optimizer with optional LLRD
    if args.lrd < 1.0 and not args.head_only:
        param_groups = param_groups_lrd(
            model, base_lr=args.lr, weight_decay=args.weight_decay, lrd=args.lrd
        )
        optimizer = optim.AdamW(param_groups, betas=(0.9, 0.999))
    else:
        optimizer = optim.AdamW(
            [p for p in model.parameters() if p.requires_grad],
            lr=args.lr,
            weight_decay=args.weight_decay,
            betas=(0.9, 0.999),
        )

    scheduler = CosineAnnealingLR(optimizer, T_max=args.epochs)
    scaler = GradScaler("cuda") if args.amp else None

    best_val_acc = 0.0
    training_history = {
        "epoch": [],
        "train_loss": [],
        "val_loss": [],
        "val_acc": [],
        "lr": [],
    }

    logger.info("Starting training...")
    start_time = time.time()

    # Create overall training progress bar
    epoch_pbar = tqdm(range(1, args.epochs + 1), desc="Training", unit="epoch")

    def maybe_augment(x, y):
        # Simple mixup/cutmix implementation (mutually exclusive here for simplicity)
        if args.mixup > 0 and not is_multilabel:
            lam = numpy_beta(args.mixup, args.mixup)
            index = torch.randperm(x.size(0), device=x.device)
            x = lam * x + (1 - lam) * x[index, :]
            y = lam * y + (1 - lam) * y[index]
            return x, y
        # (CutMix omitted for brevity; typically requires bounding-box ops)
        return x, y

    for epoch in epoch_pbar:
        model.train()
        running_loss = 0.0
        seen = 0

        # Create progress bar for training
        train_pbar = tqdm(
            train_loader, desc=f"Epoch {epoch}/{args.epochs}", leave=False, unit="batch"
        )

        for xb, yb in train_pbar:
            xb = xb.to(device, non_blocking=True)
            yb = yb.to(device, non_blocking=True)

            if not is_multilabel and args.mixup > 0:
                yb_onehot = nn.functional.one_hot(yb, num_classes=n_classes).float()
                xb, targets = maybe_augment(xb, yb_onehot)
                loss_fn = nn.BCEWithLogitsLoss()
            else:
                targets = yb
                loss_fn = criterion

            optimizer.zero_grad(set_to_none=True)
            if args.amp:
                with autocast("cuda"):
                    logits = model(xb)
                    loss = loss_fn(logits, targets)
            else:
                logits = model(xb)
                loss = loss_fn(logits, targets)
            if args.amp:
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                loss.backward()
                optimizer.step()

            running_loss += loss.item() * xb.size(0)
            seen += xb.size(0)

            # Update progress bar with current loss
            current_loss = running_loss / max(seen, 1)
            train_pbar.set_postfix({"loss": f"{current_loss:.4f}"})

        scheduler.step()

        val_loss, val_acc = evaluate(
            model, val_loader, device, is_multilabel, "Validating"
        )
        train_loss = running_loss / max(seen, 1)
        current_lr = optimizer.param_groups[0]["lr"]

        # Log metrics and update history
        training_history = log_training_metrics(
            epoch, train_loss, val_loss, val_acc, current_lr, logger, training_history
        )

        # Update epoch progress bar with validation metrics
        epoch_pbar.set_postfix(
            {
                "val_acc": f"{val_acc:.3f}",
                "val_loss": f"{val_loss:.4f}",
                "lr": f"{current_lr:.2e}",
            }
        )

        # Save checkpoint
        is_best = val_acc > best_val_acc
        if is_best:
            best_val_acc = val_acc

        save_checkpoint(
            model,
            optimizer,
            scheduler,
            epoch,
            val_acc,
            train_loss,
            val_loss,
            output_dir,
            is_best=is_best,
            logger=logger,
        )

    # Training completed
    total_time = time.time() - start_time
    logger.info(f"Training completed in {total_time/3600:.2f} hours")

    # Final test
    # Load best checkpoint
    best_checkpoint_path = output_dir / "checkpoints" / "best.pth"
    if best_checkpoint_path.exists():
        ckpt = torch.load(best_checkpoint_path, map_location=device, weights_only=False)
        # Handle different checkpoint key formats
        if "model_state_dict" in ckpt:
            model.load_state_dict(ckpt["model_state_dict"])
        elif "model" in ckpt:
            model.load_state_dict(ckpt["model"])
        else:
            logger.warning(
                "Could not find model weights in checkpoint, using current model"
            )
            ckpt = {}  # Set to empty dict to avoid errors below

        logger.info(
            f"Loaded best checkpoint from epoch {ckpt.get('epoch','?')} "
            f"with val_acc={ckpt.get('val_acc',0)*100:.2f}%"
        )
    else:
        logger.warning("No best checkpoint found, using current model")

    test_loss, test_acc = evaluate(model, test_loader, device, is_multilabel, "Testing")
    logger.info(f"FINAL TEST: loss={test_loss:.4f} | acc={test_acc*100:.2f}%")

    # Save training summary
    summary = {
        "final_test_acc": test_acc,
        "final_test_loss": test_loss,
        "best_val_acc": best_val_acc,
        "total_training_time_hours": total_time / 3600,
        "total_epochs": args.epochs,
        "training_history": training_history,
        "args": vars(args),
    }

    summary_path = output_dir / "logs" / "training_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)

    logger.info(f"Training summary saved to {summary_path}")
    logger.info(f"All outputs saved to {output_dir}")


def numpy_beta(a, b):
    # small helper to sample from Beta(a,b) without numpy dependency
    # using torch distributions
    dist = torch.distributions.Beta(torch.tensor([a]), torch.tensor([b]))
    return float(dist.sample().item())


if __name__ == "__main__":
    main()