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)