train_ssv2_classifier.py

#!/usr/bin/env python3
"""
Production training script for SSv2 video classifier with frozen V-JEPA 2 encoder.

This script trains an attentive classifier head on top of a frozen pretrained V-JEPA 2
encoder using the Something-Something-v2 dataset with MLX on Apple Silicon.
"""

import argparse
import json
import logging
import sys
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Tuple, Optional

import cv2
import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
from mlx.utils import tree_flatten
import numpy as np
from tqdm import tqdm

# Add project root to path
PROJECT_ROOT = Path(__file__).parent
sys.path.insert(0, str(PROJECT_ROOT))

from vjepa2_mlx.models.vision_transformer import VisionTransformer
from vjepa2_mlx.models.attentive_pooler import AttentiveClassifier

# Constants
IMAGENET_DEFAULT_MEAN = np.array([0.485, 0.456, 0.406])
IMAGENET_DEFAULT_STD = np.array([0.229, 0.224, 0.225])


def setup_logging(output_dir: Path, verbose: bool = False) -> logging.Logger:
    """Setup logging configuration."""
    log_level = logging.DEBUG if verbose else logging.INFO
    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    
    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Setup file handler
    log_file = output_dir / f'training_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'
    
    logging.basicConfig(
        level=log_level,
        format=log_format,
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )
    
    logger = logging.getLogger(__name__)
    logger.info(f"Logging to {log_file}")
    return logger


def load_video(video_path: Path, num_frames: int = 16, target_size: Tuple[int, int] = (224, 224)) -> Optional[np.ndarray]:
    """Load video and extract frames uniformly."""
    try:
        cap = cv2.VideoCapture(str(video_path))
        frames = []
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            # Convert BGR to RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frames.append(frame)
        
        cap.release()
        
        if len(frames) == 0:
            return None
        
        # Sample frames uniformly
        if len(frames) >= num_frames:
            indices = np.linspace(0, len(frames) - 1, num_frames, dtype=int)
            frames = [frames[i] for i in indices]
        else:
            # Repeat last frame if not enough frames
            while len(frames) < num_frames:
                frames.append(frames[-1])
        
        # Resize frames
        frames = [cv2.resize(f, target_size) for f in frames]
        
        return np.stack(frames)  # (T, H, W, C)
    except Exception as e:
        logging.error(f"Error loading video {video_path}: {e}")
        return None


def preprocess_video(frames: np.ndarray, mean: np.ndarray = IMAGENET_DEFAULT_MEAN, 
                     std: np.ndarray = IMAGENET_DEFAULT_STD) -> mx.array:
    """Preprocess video frames for model input."""
    # Normalize to [0, 1]
    frames = frames.astype(np.float32) / 255.0
    
    # Apply ImageNet normalization
    frames = (frames - mean) / std
    
    # Convert to MLX array (keep as T, H, W, C for MLX Conv3d)
    return mx.array(frames)


def augment_video(frames: np.ndarray, training: bool = True) -> np.ndarray:
    """Apply data augmentation to video frames."""
    if not training:
        return frames
    
    # Random horizontal flip
    if np.random.rand() > 0.5:
        frames = np.flip(frames, axis=2)  # Flip width dimension
    
    # Random crop
    if np.random.rand() > 0.5:
        h, w = frames.shape[1:3]
        crop_size = int(0.9 * min(h, w))
        top = np.random.randint(0, h - crop_size + 1)
        left = np.random.randint(0, w - crop_size + 1)
        frames = frames[:, top:top+crop_size, left:left+crop_size, :]
        # Resize back
        frames = np.stack([cv2.resize(f, (w, h)) for f in frames])
    
    return frames


class SSv2Dataset:
    """Something-Something-v2 dataset loader."""
    
    def __init__(self, data: List[Dict], videos_dir: Path, label_to_idx: Dict[str, int],
                 num_frames: int = 16, resolution: int = 224, training: bool = True,
                 logger: Optional[logging.Logger] = None):
        self.data = data
        self.videos_dir = Path(videos_dir)
        self.label_to_idx = label_to_idx
        self.num_frames = num_frames
        self.resolution = resolution
        self.training = training
        self.logger = logger or logging.getLogger(__name__)
        
        # Filter samples with existing videos
        self.valid_samples = []
        for sample in data:
            video_path = self.videos_dir / f"{sample['id']}.webm"
            if video_path.exists():
                self.valid_samples.append(sample)
        
        self.logger.info(f"Valid samples: {len(self.valid_samples)} / {len(data)}")
    
    def __len__(self) -> int:
        return len(self.valid_samples)
    
    def __getitem__(self, idx: int) -> Tuple[mx.array, int]:
        sample = self.valid_samples[idx]
        video_path = self.videos_dir / f"{sample['id']}.webm"
        
        # Load video
        frames = load_video(video_path, self.num_frames, (self.resolution, self.resolution))
        
        if frames is None:
            raise ValueError(f"Failed to load video: {video_path}")
        
        # Apply augmentation
        frames = augment_video(frames, training=self.training)
        
        # Preprocess
        frames = preprocess_video(frames)
        
        # Get label index (strip brackets from template)
        template = sample['template'].replace('[', '').replace(']', '')
        label_idx = int(self.label_to_idx[template])
        
        return frames, label_idx


def cross_entropy_loss(logits: mx.array, targets: mx.array, num_classes: int) -> mx.array:
    """Compute cross entropy loss."""
    # Convert targets to one-hot
    targets_one_hot = mx.zeros((targets.shape[0], num_classes))
    targets_one_hot[mx.arange(targets.shape[0]), targets] = 1
    
    # Compute log softmax
    log_probs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
    
    # Compute loss
    loss = -mx.sum(targets_one_hot * log_probs) / targets.shape[0]
    return loss


def accuracy(logits: mx.array, targets: mx.array) -> mx.array:
    """Compute top-1 accuracy."""
    predictions = mx.argmax(logits, axis=-1)
    return mx.mean(predictions == targets)


def top_k_accuracy(logits: mx.array, targets: mx.array, k: int = 5) -> mx.array:
    """Compute top-k accuracy."""
    top_k_preds = mx.argsort(logits, axis=-1)[:, -k:]
    targets_expanded = mx.expand_dims(targets, axis=-1)
    correct = mx.any(top_k_preds == targets_expanded, axis=-1)
    return mx.mean(correct.astype(mx.float32))


class CosineSchedule:
    """Cosine learning rate schedule with warmup."""
    
    def __init__(self, base_lr: float, final_lr: float, total_steps: int, warmup_steps: int):
        self.base_lr = base_lr
        self.final_lr = final_lr
        self.total_steps = total_steps
        self.warmup_steps = warmup_steps
    
    def __call__(self, step: int) -> float:
        if step < self.warmup_steps:
            # Linear warmup
            return self.base_lr * (step / self.warmup_steps)
        else:
            # Cosine decay
            progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            return self.final_lr + (self.base_lr - self.final_lr) * 0.5 * (1 + np.cos(np.pi * progress))


def forward_pass(encoder: nn.Module, classifier: nn.Module, video: mx.array, frozen_encoder: bool = True) -> mx.array:
    """Forward pass through encoder and classifier."""
    if frozen_encoder:
        features = encoder(video)
        features = mx.stop_gradient(features)
    else:
        features = encoder(video)
    
    logits = classifier(features)
    return logits


def loss_fn(encoder: nn.Module, classifier: nn.Module, video: mx.array, targets: mx.array, num_classes: int) -> Tuple[mx.array, mx.array]:
    """Compute loss for a batch."""
    logits = forward_pass(encoder, classifier, video, frozen_encoder=True)
    loss = cross_entropy_loss(logits, targets, num_classes)
    return loss, logits


def scale_gradients(grads, scale: float):
    """Recursively scale gradients by a factor."""
    if isinstance(grads, dict):
        return {k: scale_gradients(v, scale) for k, v in grads.items()}
    elif isinstance(grads, (list, tuple)):
        return type(grads)(scale_gradients(v, scale) for v in grads)
    else:
        # Assume it's an array
        return grads * scale


def add_gradients(grads1, grads2):
    """Recursively add two gradient structures."""
    if isinstance(grads1, dict):
        return {k: add_gradients(grads1[k], grads2[k]) for k in grads1.keys()}
    elif isinstance(grads1, (list, tuple)):
        return type(grads1)(add_gradients(g1, g2) for g1, g2 in zip(grads1, grads2))
    else:
        # Assume it's an array
        return grads1 + grads2


def train_step(encoder: nn.Module, classifier: nn.Module, optimizer: optim.Optimizer,
               video: mx.array, targets: mx.array, loss_and_grad_fn, num_classes: int,
               accumulate_grads: bool = False) -> Tuple[float, float, Optional[Dict]]:
    """Single training step with optional gradient accumulation."""
    # Compute loss and gradients
    (loss, logits), grads = loss_and_grad_fn(encoder, classifier, video, targets, num_classes)
    
    # Update classifier parameters only if not accumulating
    if not accumulate_grads:
        optimizer.update(classifier, grads)
    
    # Compute metrics
    acc = accuracy(logits, targets)
    
    return float(loss), float(acc), grads if accumulate_grads else None


def eval_step(encoder: nn.Module, classifier: nn.Module, video: mx.array, targets: mx.array, num_classes: int) -> Tuple[float, float, float]:
    """Single evaluation step."""
    logits = forward_pass(encoder, classifier, video, frozen_encoder=True)
    loss = cross_entropy_loss(logits, targets, num_classes)
    acc = accuracy(logits, targets)
    top5_acc = top_k_accuracy(logits, targets, k=5)
    return float(loss), float(acc), float(top5_acc)


def train_epoch(encoder: nn.Module, classifier: nn.Module, optimizer: optim.Optimizer,
                train_dataset: SSv2Dataset, batch_size: int, lr_schedule: CosineSchedule,
                loss_and_grad_fn, num_classes: int, global_step: int, logger: logging.Logger,
                use_wandb: bool = False, gradient_accumulation_steps: int = 1, 
                save_every_steps: int = 1000, output_dir: Optional[Path] = None) -> Tuple[float, float, int]:
    """Train for one epoch with gradient accumulation."""
    classifier.train()
    train_losses = []
    train_accs = []
    
    # Shuffle training data
    indices = np.random.permutation(len(train_dataset))
    
    # For gradient accumulation
    accumulated_grads = None
    accumulation_count = 0
    
    for batch_idx in tqdm(range(0, len(train_dataset), batch_size), desc="Training"):
        # Get batch indices
        batch_indices = indices[batch_idx:batch_idx + batch_size]
        if len(batch_indices) < batch_size:
            continue  # Skip incomplete batch
        
        # Load batch
        videos = []
        targets = []
        for idx in batch_indices:
            try:
                video, target = train_dataset[int(idx)]
                videos.append(video)
                targets.append(target)
            except Exception as e:
                logger.warning(f"Failed to load sample {idx}: {e}")
                continue
        
        if len(videos) == 0:
            continue
        
        videos = mx.stack(videos)
        targets = mx.array(targets)
        
        # Update learning rate
        lr = lr_schedule(global_step)
        optimizer.learning_rate = lr
        
        # Training step with gradient accumulation
        should_accumulate = (accumulation_count + 1) < gradient_accumulation_steps
        loss, acc, grads = train_step(encoder, classifier, optimizer, videos, targets, 
                                      loss_and_grad_fn, num_classes, accumulate_grads=should_accumulate)
        
        # Accumulate gradients
        if should_accumulate:
            if accumulated_grads is None:
                accumulated_grads = grads
            else:
                # Add current gradients to accumulated
                accumulated_grads = add_gradients(accumulated_grads, grads)
            
            accumulation_count += 1
        else:
            # Apply accumulated gradients
            if accumulated_grads is not None:
                # Average the gradients
                scale_factor = 1.0 / gradient_accumulation_steps
                accumulated_grads = scale_gradients(accumulated_grads, scale_factor)
                
                # Update with accumulated gradients
                optimizer.update(classifier, accumulated_grads)
                mx.eval(optimizer.state)
            
            # Reset accumulation
            accumulated_grads = None
            accumulation_count = 0
            
            # Only count as a step when we actually update
            global_step += 1
        
        train_losses.append(loss)
        train_accs.append(acc)
        
        # Log to wandb (only on actual updates)
        if use_wandb and not should_accumulate:
            import wandb
            wandb.log({
                "train/loss": loss,
                "train/acc": acc,
                "train/learning_rate": lr,
                "train/step": global_step,
            })
        
        # Log progress
        if not should_accumulate and (global_step % 10 == 0):
            logger.info(f"  Step {global_step}: loss={loss:.4f}, acc={acc:.4f}, lr={lr:.6f}")
        
        # Save checkpoint every N steps
        if not should_accumulate and output_dir is not None and (global_step % save_every_steps == 0):
            checkpoint_path = output_dir / f"classifier_step_{global_step}.safetensors"
            try:
                # Get current epoch from function attribute
                current_epoch = getattr(train_epoch, '_current_epoch', 0)
                save_checkpoint(classifier, checkpoint_path, logger, 
                              epoch=current_epoch, global_step=global_step)
            except Exception as e:
                logger.error(f"Failed to save step checkpoint: {e}")
    
    epoch_train_loss = float(np.mean(train_losses))
    epoch_train_acc = float(np.mean(train_accs))
    
    return epoch_train_loss, epoch_train_acc, global_step


def load_checkpoint(classifier: nn.Module, checkpoint_path: Path, logger: logging.Logger) -> Dict:
    """Load model checkpoint and return training state."""
    try:
        logger.info(f"Loading checkpoint from {checkpoint_path}")
        weights, metadata = mx.load(str(checkpoint_path), return_metadata=True)
        
        # Load weights into classifier
        classifier.load_weights(list(weights.items()))
        # Load metadata into training state
        
        # Try to extract training state from filename
        training_state = metadata
        
        logger.info(f"Checkpoint loaded successfully")
        return training_state
        
    except Exception as e:
        logger.error(f"Failed to load checkpoint: {e}")
        raise


def validate(encoder: nn.Module, classifier: nn.Module, val_dataset: SSv2Dataset,
             batch_size: int, num_classes: int, logger: logging.Logger) -> Tuple[float, float, float]:
    """Validate the model."""
    classifier.eval()
    val_losses = []
    val_accs = []
    val_top5_accs = []
    
    for batch_idx in tqdm(range(0, len(val_dataset), batch_size), desc="Validation"):
        batch_end = min(batch_idx + batch_size, len(val_dataset))
        
        videos = []
        targets = []
        for idx in range(batch_idx, batch_end):
            try:
                video, target = val_dataset[idx]
                videos.append(video)
                targets.append(target)
            except Exception as e:
                logger.warning(f"Failed to load validation sample {idx}: {e}")
                continue
        
        if len(videos) == 0:
            continue
        
        videos = mx.stack(videos)
        targets = mx.array(targets)
        
        # Evaluation step
        loss, acc, top5_acc = eval_step(encoder, classifier, videos, targets, num_classes)
        mx.eval(loss, acc, top5_acc)
        
        val_losses.append(loss)
        val_accs.append(acc)
        val_top5_accs.append(top5_acc)
    
    epoch_val_loss = np.mean(val_losses)
    epoch_val_acc = np.mean(val_accs)
    epoch_val_top5_acc = np.mean(val_top5_accs)
    
    return epoch_val_loss, epoch_val_acc, epoch_val_top5_acc


def save_checkpoint(classifier: nn.Module, checkpoint_path: Path, logger: logging.Logger,
                   epoch: Optional[int] = None, global_step: Optional[int] = None,
                   best_val_acc: Optional[float] = None):
    """Save model checkpoint with training state metadata."""
    try:
        # Create metadata
        metadata = {}
        if epoch is not None:
            metadata['epoch'] = str(epoch)
        if global_step is not None:
            metadata['global_step'] = str(global_step)
        if best_val_acc is not None:
            metadata['best_val_acc'] = str(best_val_acc)

        flat_params: dict[str, mx.array] = tree_flatten(classifier.parameters(), destination={})

        mx.save_safetensors(str(checkpoint_path), flat_params, metadata=metadata)
        
        logger.info(f"Saved checkpoint to {checkpoint_path}")
        if metadata:
            logger.info(f"  Metadata: {metadata}")
            
    except Exception as e:
        logger.error(f"Failed to save checkpoint: {e}")


def main():
    parser = argparse.ArgumentParser(description="Train SSv2 video classifier with frozen V-JEPA 2 encoder")
    
    # Dataset arguments
    parser.add_argument("--videos-dir", type=Path, required=True, help="Path to videos directory")
    parser.add_argument("--labels-dir", type=Path, required=True, help="Path to labels directory")
    parser.add_argument("--train-json", type=str, default="train.json", help="Training data JSON filename")
    parser.add_argument("--val-json", type=str, default="validation.json", help="Validation data JSON filename")
    parser.add_argument("--subset-size", type=int, default=None, help="Use subset of data for faster training")
    
    # Model arguments
    parser.add_argument("--pretrained-weights", type=Path, required=True, help="Path to pretrained encoder weights")
    parser.add_argument("--num-classes", type=int, default=174, help="Number of classes")
    parser.add_argument("--frames-per-clip", type=int, default=16, help="Number of frames per video clip")
    parser.add_argument("--resolution", type=int, default=224, help="Video frame resolution")
    parser.add_argument("--tubelet-size", type=int, default=2, help="Tubelet size for 3D patch embedding")
    parser.add_argument("--num-probe-blocks", type=int, default=1, help="Number of probe blocks in classifier")
    parser.add_argument("--num-heads", type=int, default=16, help="Number of attention heads in classifier")
    
    # Training arguments
    parser.add_argument("--batch-size", type=int, default=4, help="Batch size")
    parser.add_argument("--num-epochs", type=int, default=10, help="Number of training epochs")
    parser.add_argument("--learning-rate", type=float, default=1e-3, help="Learning rate")
    parser.add_argument("--weight-decay", type=float, default=0.05, help="Weight decay")
    parser.add_argument("--warmup-epochs", type=int, default=1, help="Number of warmup epochs")
    parser.add_argument("--gradient-accumulation-steps", type=int, default=1, help="Number of gradient accumulation steps")
    
    # Weights & Biases arguments
    parser.add_argument("--use-wandb", action="store_true", help="Enable Weights & Biases logging")
    parser.add_argument("--wandb-project", type=str, default="vjepa2-ssv2-classifier", help="W&B project name")
    parser.add_argument("--wandb-entity", type=str, default=None, help="W&B entity (username/team)")
    parser.add_argument("--wandb-run-name", type=str, default=None, help="W&B run name")
    
    # Output arguments
    parser.add_argument("--output-dir", type=Path, default=Path("output_ssv2_classifier"), help="Output directory")
    parser.add_argument("--save-every-steps", type=int, default=1000, help="Save checkpoint every N steps")
    parser.add_argument("--resume-from", type=Path, default=None, help="Resume training from checkpoint path")
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging")
    
    args = parser.parse_args()
    
    # Setup logging
    logger = setup_logging(args.output_dir, args.verbose)
    logger.info(f"MLX device: {mx.default_device()}")
    logger.info(f"Arguments: {vars(args)}")
    
    # Load labels
    logger.info("Loading labels...")
    with open(args.labels_dir / "labels.json", "r") as f:
        label_to_idx = json.load(f)
    
    idx_to_label = {int(v): k for k, v in label_to_idx.items()}
    logger.info(f"Number of classes: {len(label_to_idx)}")
    
    # Load training and validation data
    logger.info("Loading dataset metadata...")
    with open(args.labels_dir / args.train_json, "r") as f:
        train_data = json.load(f)
    
    with open(args.labels_dir / args.val_json, "r") as f:
        val_data = json.load(f)
    
    # Apply subset if specified
    if args.subset_size is not None:
        logger.info(f"Using subset of {args.subset_size} training samples")
        train_data = train_data[:args.subset_size]
        val_data = val_data[:args.subset_size // 5]
    
    logger.info(f"Training samples: {len(train_data)}")
    logger.info(f"Validation samples: {len(val_data)}")
    
    # Create datasets
    logger.info("Creating datasets...")
    train_dataset = SSv2Dataset(
        train_data, args.videos_dir, label_to_idx,
        num_frames=args.frames_per_clip, resolution=args.resolution,
        training=True, logger=logger
    )
    
    val_dataset = SSv2Dataset(
        val_data, args.videos_dir, label_to_idx,
        num_frames=args.frames_per_clip, resolution=args.resolution,
        training=False, logger=logger
    )
    
    # Initialize encoder
    logger.info("Initializing encoder...")
    encoder = VisionTransformer(
        img_size=(args.resolution, args.resolution),
        patch_size=16,
        num_frames=args.frames_per_clip,
        tubelet_size=args.tubelet_size,
        in_chans=3,
        embed_dim=1024,
        depth=24,
        num_heads=16,
        mlp_ratio=4.0,
        qkv_bias=True,
        use_silu=False,
        wide_silu=False,
        use_rope=True,
    )
    
    # Load pretrained weights
    if args.pretrained_weights.exists():
        logger.info(f"Loading pretrained weights from {args.pretrained_weights}")
        weights = mx.load(str(args.pretrained_weights))
        encoder.load_weights(list(weights.items()))
        logger.info("Pretrained weights loaded successfully")
    else:
        logger.warning(f"Pretrained weights not found at {args.pretrained_weights}")
        logger.warning("Training will start from random initialization")
    
    # Initialize classifier
    logger.info("Initializing classifier...")
    classifier = AttentiveClassifier(
        embed_dim=encoder.embed_dim,
        num_heads=args.num_heads,
        mlp_ratio=4.0,
        depth=args.num_probe_blocks,
        num_classes=args.num_classes,
        qkv_bias=True,
        complete_block=True,
    )
    
    # Resume from checkpoint if specified
    resume_epoch = 0
    resume_step = 0
    resume_best_val_acc = 0.0
    
    if args.resume_from is not None:
        if args.resume_from.exists():
            logger.info(f"Resuming from checkpoint: {args.resume_from}")
            training_state = load_checkpoint(classifier, args.resume_from, logger)
            # Convert string values to proper types
            resume_epoch = int(training_state.get('epoch', 0))
            resume_step = int(training_state.get('global_step', 0))
            resume_best_val_acc = float(training_state.get('best_val_acc', 0.0))
            logger.info(f"Resuming from epoch {resume_epoch}, step {resume_step}")
        else:
            logger.warning(f"Resume checkpoint not found: {args.resume_from}")
            logger.warning("Starting training from scratch")
    
    # Setup training
    steps_per_epoch = len(train_dataset) // args.batch_size
    total_steps = steps_per_epoch * args.num_epochs
    warmup_steps = steps_per_epoch * args.warmup_epochs
    
    logger.info(f"Steps per epoch: {steps_per_epoch}")
    logger.info(f"Total steps: {total_steps}")
    logger.info(f"Warmup steps: {warmup_steps}")
    
    # Create learning rate schedule
    lr_schedule = CosineSchedule(
        base_lr=args.learning_rate,
        final_lr=args.learning_rate * 0.01,
        total_steps=total_steps,
        warmup_steps=warmup_steps
    )
    
    # Initialize optimizer
    optimizer = optim.AdamW(learning_rate=args.learning_rate, weight_decay=args.weight_decay)
    
    # Create loss and grad function
    loss_and_grad_fn = nn.value_and_grad(classifier, loss_fn)
    
    # Initialize wandb
    if args.use_wandb:
        import wandb
        run_name = args.wandb_run_name or f"ssv2_vitl_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        wandb.init(
            project=args.wandb_project,
            entity=args.wandb_entity,
            name=run_name,
            config=vars(args)
        )
        logger.info(f"Wandb initialized: {run_name}")
    
    # Training loop
    logger.info("Starting training...")
    history = {
        'train_loss': [],
        'train_acc': [],
        'val_loss': [],
        'val_acc': [],
        'val_top5_acc': [],
    }
    
    best_val_acc = resume_best_val_acc
    global_step = resume_step
    
    if resume_epoch > 0:
        logger.info(f"Resuming training from epoch {resume_epoch + 1}")
        logger.info(f"Best validation accuracy so far: {best_val_acc:.4f}")
    
    for epoch in range(resume_epoch, args.num_epochs):
        logger.info(f"\n{'='*60}")
        logger.info(f"Epoch {epoch + 1}/{args.num_epochs}")
        logger.info(f"{'='*60}")
        
        # Training phase
        # Store current epoch for checkpoint saving
        train_epoch._current_epoch = epoch
        epoch_train_loss, epoch_train_acc, global_step = train_epoch(
            encoder, classifier, optimizer, train_dataset, args.batch_size,
            lr_schedule, loss_and_grad_fn, args.num_classes, global_step, logger, 
            args.use_wandb, args.gradient_accumulation_steps, args.save_every_steps, args.output_dir
        )
        
        history['train_loss'].append(epoch_train_loss)
        history['train_acc'].append(epoch_train_acc)
        
        logger.info(f"Training - Loss: {epoch_train_loss:.4f}, Acc: {epoch_train_acc:.4f}")
        
        # Log epoch metrics to wandb
        if args.use_wandb:
            import wandb
            wandb.log({
                "epoch": epoch + 1,
                "train/epoch_loss": epoch_train_loss,
                "train/epoch_acc": epoch_train_acc,
            })
        
        # Validation phase
        epoch_val_loss, epoch_val_acc, epoch_val_top5_acc = validate(
            encoder, classifier, val_dataset, args.batch_size, args.num_classes, logger
        )
        
        history['val_loss'].append(epoch_val_loss)
        history['val_acc'].append(epoch_val_acc)
        history['val_top5_acc'].append(epoch_val_top5_acc)
        
        logger.info(f"Validation - Loss: {epoch_val_loss:.4f}, Acc: {epoch_val_acc:.4f}, Top-5 Acc: {epoch_val_top5_acc:.4f}")
        
        # Log validation metrics to wandb
        if args.use_wandb:
            import wandb
            wandb.log({
                "val/epoch_loss": epoch_val_loss,
                "val/epoch_acc": epoch_val_acc,
                "val/epoch_top5_acc": epoch_val_top5_acc,
            })
        
        # Save best model
        if epoch_val_acc > best_val_acc:
            best_val_acc = epoch_val_acc
            checkpoint_path = args.output_dir / "best_classifier.safetensors"
            try:
                save_checkpoint(classifier, checkpoint_path, logger)
                logger.info(f"✓ New best model! Val acc: {best_val_acc:.4f}")
            except Exception as e:
                logger.error(f"Failed to save best model: {e}")
            
            if args.use_wandb:
                import wandb
                if wandb.run is not None:
                    wandb.run.summary["best_val_acc"] = best_val_acc
                    wandb.run.summary["best_epoch"] = epoch + 1
    
    # Save final model and history
    logger.info("\n" + "="*60)
    logger.info("Training completed!")
    logger.info(f"Best validation accuracy: {best_val_acc:.4f}")
    logger.info("="*60)
    
    # Save training history
    history_path = args.output_dir / 'training_history.json'
    with open(history_path, 'w') as f:
        json.dump(history, f, indent=2)
    logger.info(f"Training history saved to {history_path}")
    
    # Finish wandb
    if args.use_wandb:
        import wandb
        wandb.finish()
        logger.info("Wandb run finished")


if __name__ == "__main__":
    main()