scripts/train_classifier.py

import os
import sys
import argparse
import time
import yaml
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

# Ensure repo root is on sys.path (fixes ModuleNotFoundError for src/)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.data_loader import get_eurosat_dataloaders
from src.cnn_model import GreeneryClassifier


def load_config(config_path="config/config.yaml"):
    with open(config_path, "r") as f:
        return yaml.safe_load(f)


def gpu_handshake():
    """
    Strict GPU verification. Prints the GPU name on success.
    Raises RuntimeError and halts execution if no CUDA GPU is found.
    """
    if not torch.cuda.is_available():
        raise RuntimeError(
            "❌ FATAL: No CUDA-capable GPU detected!\n"
            "   This training script requires an NVIDIA GPU with CUDA support.\n"
            "   Please verify:\n"
            "     1. Your NVIDIA drivers are installed (nvidia-smi)\n"
            "     2. You installed the CUDA version of PyTorch (torch+cu...)\n"
            "     3. Your GPU is visible to the system\n"
            "   Aborting to prevent silent CPU fallback."
        )

    gpu_name = torch.cuda.get_device_name(0)
    vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    print(f"🚀 Training on: {gpu_name} ({vram_gb:.1f} GB VRAM)")
    print(f"   CUDA Version: {torch.version.cuda}")
    print(f"   PyTorch Version: {torch.__version__}")
    return torch.device("cuda")


def train(config, args):
    device = gpu_handshake()

    train_loader, val_loader, classes = get_eurosat_dataloaders(
        data_dir=config["paths"]["eurosat_dir"],
        batch_size=config["training"]["batch_size"],
    )
    print(
        f"📊 Dataset loaded: {len(train_loader.dataset)} train / {len(val_loader.dataset)} val samples"
    )

    model = GreeneryClassifier(
        num_classes=config["model"]["num_classes"], pretrained=True
    )
    model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])

    # TensorBoard Setup
    log_dir = config["paths"].get("output_logs", "outputs/logs")
    os.makedirs(log_dir, exist_ok=True)
    writer = SummaryWriter(log_dir=log_dir)
    print(f"📈 TensorBoard logs → {os.path.abspath(log_dir)}")

    # AMP Setup (Automatic Mixed Precision — leverages Tensor Cores on RTX cards)
    scaler = torch.amp.GradScaler("cuda")

    best_val_acc = 0.0
    start_epoch = 0
    checkpoint_path = os.path.join(
        config["paths"]["output_models"], "training_checkpoint.pth"
    )
    best_model_path = os.path.join(
        config["paths"]["output_models"], "resnet50_eurosat.pth"
    )

    # Resume from checkpoint if enabled and a checkpoint exists
    if config["training"].get("resume_checkpoint", False) and os.path.exists(
        checkpoint_path
    ):
        print(f"🔄 Resuming from checkpoint: {checkpoint_path}")
        checkpoint = torch.load(checkpoint_path, map_location=device)
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        scaler.load_state_dict(checkpoint["scaler_state_dict"])
        start_epoch = checkpoint["epoch"] + 1
        best_val_acc = checkpoint["best_val_acc"]
        print(
            f"   ↳ Resumed at epoch {start_epoch}/{config['training']['epochs']} | Best Val Acc so far: {best_val_acc:.2f}%"
        )

    if args.dry_run:
        print("✅ Dry run completed successfully. Models and DataLoaders initialized.")
        writer.close()
        return

    epochs = config["training"]["epochs"]
    training_start = time.time()
    print(f"\n{'='*60}")
    print(
        f"  TRAINING START — {epochs} epochs, batch_size={config['training']['batch_size']}"
    )
    print(
        f"  AMP: Enabled | Optimizer: Adam | LR: {config['training']['learning_rate']}"
    )
    print(f"{'='*60}\n")

    for epoch in range(start_epoch, epochs):
        epoch_start = time.time()
        model.train()
        running_loss = 0.0

        # Training loop with tqdm progress bar
        train_bar = tqdm(
            train_loader,
            desc=f"Epoch {epoch+1}/{epochs} [Train]",
            unit="batch",
            leave=True,
            ncols=100,
        )
        for inputs, labels in train_bar:
            inputs, labels = inputs.to(device, non_blocking=True), labels.to(
                device, non_blocking=True
            )
            optimizer.zero_grad()

            # AMP Autocast — forward pass in float16 for speed
            with torch.autocast(device_type="cuda"):
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            running_loss += loss.item()
            train_bar.set_postfix(loss=f"{loss.item():.4f}")

        # Validation loop with tqdm progress bar
        model.eval()
        correct = 0
        total = 0
        val_loss = 0.0
        with torch.no_grad():
            val_bar = tqdm(
                val_loader,
                desc=f"Epoch {epoch+1}/{epochs} [Val]  ",
                unit="batch",
                leave=True,
                ncols=100,
            )
            for inputs, labels in val_bar:
                inputs, labels = inputs.to(device, non_blocking=True), labels.to(
                    device, non_blocking=True
                )
                with torch.autocast(device_type="cuda"):
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                val_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        train_loss = running_loss / len(train_loader)
        val_loss_avg = val_loss / len(val_loader)
        val_acc = 100 * correct / total
        epoch_time = time.time() - epoch_start

        print(
            f"  ✦ Epoch {epoch+1}/{epochs} — "
            f"Train Loss: {train_loss:.4f} | Val Loss: {val_loss_avg:.4f} | "
            f"Val Acc: {val_acc:.2f}% | Time: {epoch_time:.1f}s"
        )

        # Log to TensorBoard
        writer.add_scalar("Loss/Train", train_loss, epoch)
        writer.add_scalar("Loss/Validation", val_loss_avg, epoch)
        writer.add_scalar("Accuracy/Validation", val_acc, epoch)
        writer.add_scalar("Time/Epoch_Seconds", epoch_time, epoch)

        # Save checkpoint every epoch (for crash recovery)
        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
        torch.save(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "scaler_state_dict": scaler.state_dict(),
                "best_val_acc": best_val_acc,
            },
            checkpoint_path,
        )

        # Save best model
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), best_model_path)
            print(
                f"  🏆 New best model saved! Val Acc: {val_acc:.2f}% → {best_model_path}"
            )

    total_time = time.time() - training_start
    print(f"\n{'='*60}")
    print(f"  ✅ TRAINING COMPLETE")
    print(f"  Best Validation Accuracy: {best_val_acc:.2f}%")
    print(f"  Total Training Time: {total_time/60:.1f} minutes")
    print(f"  Best Model: {os.path.abspath(best_model_path)}")
    print(f"  TensorBoard Logs: {os.path.abspath(log_dir)}")
    print(f"{'='*60}\n")
    writer.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train CNN Classifier on EuroSAT")
    parser.add_argument("--config", default="config/config.yaml", help="Path to config")
    parser.add_argument(
        "--dry-run", action="store_true", help="Initialize but do not train"
    )
    args = parser.parse_args()

    config = load_config(args.config)
    train(config, args)