train.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
from sklearn.utils import class_weight
from sklearn.metrics import classification_report
from sklearn.model_selection import train_test_split

from src.dataset import SolarDataset, train_transforms, val_transforms
from src.model import SolarResNet
from src.engine import train_one_epoch, validate

def main():
    IMAGE_SIZE = (224, 224)
    BATCH_SIZE = 32
    DATASET_DIR = r"C:\Users\Paras\Desktop\Projects\Solar-Panel-Pytorch\data"
    
    # Fallback to relative path if absolute path doesn't exist
    if not os.path.exists(DATASET_DIR):
        DATASET_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")

    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {DEVICE}")

    # Set random seed for reproducibility
    torch.manual_seed(42)
    np.random.seed(42)

    # Scan directory paths
    base_folder = datasets.ImageFolder(root=DATASET_DIR)
    classes = base_folder.classes
    num_classes = len(classes)

    image_paths = [sample[0] for sample in base_folder.samples]
    image_labels = base_folder.targets

    train_paths, val_paths, train_labels, val_labels = train_test_split(
        image_paths, image_labels, test_size=0.2, random_state=42, stratify=image_labels
    )

    # Instantiate pipelines
    train_dataset = SolarDataset(train_paths, train_labels, train_transforms)
    val_dataset = SolarDataset(val_paths, val_labels, val_transforms)

    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0, pin_memory=True)

    print(f"Successfully loaded {len(train_dataset)} training images and {len(val_dataset)} validation images.")

    # Initialize model
    model = SolarResNet(num_classes=num_classes).to(DEVICE)

    # Compute balanced class weights
    weights = class_weight.compute_class_weight(
        class_weight='balanced',
        classes=np.unique(train_labels),
        y=train_labels
    )
    class_weights_tensor = torch.tensor(weights, dtype=torch.float32).to(DEVICE)

    criterion = nn.CrossEntropyLoss(weight=class_weights_tensor)
    
    # --- Phase 1: Training classifier head ---
    PHASE1_EPOCHS = 20
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=1e-3)

    print("\n--- Starting Training Optimization Loop (Phase 1) ---")
    for epoch in range(PHASE1_EPOCHS):
        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, DEVICE)
        val_loss, val_acc = validate(model, val_loader, criterion, DEVICE)
        
        print(f"Epoch {epoch+1}/{PHASE1_EPOCHS} -> "
              f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc*100:.2f}% || "
              f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc*100:.2f}%")

    # --- Phase 2: Fine-tuning layer4 ---
    PHASE2_EPOCHS = 15
    PATIENCE = 5

    # Unfreeze layer4 parameters
    for param in model.base_model.layer4.parameters():
        param.requires_grad = True

    optimizer = optim.Adam([
        {'params': model.base_model.layer4.parameters(), 'lr': 1e-5},
        {'params': model.base_model.fc.parameters(), 'lr': 1e-4}
    ])

    still_frozen = nn.Sequential(
        model.base_model.conv1, model.base_model.bn1,
        model.base_model.layer1, model.base_model.layer2, model.base_model.layer3
    )

    best_val_loss = float('inf')
    epochs_no_improve = 0
    best_state = None

    print("\n--- Phase 2: Fine-tuning layer4 ---")
    for epoch in range(PHASE2_EPOCHS):
        train_loss, train_acc = train_one_epoch(model, train_loader, criterion, optimizer, DEVICE,
                                                 frozen_submodule=still_frozen)
        val_loss, val_acc = validate(model, val_loader, criterion, DEVICE)
        print(f"Epoch {epoch+1}/{PHASE2_EPOCHS} -> Train Loss: {train_loss:.4f} | Train Acc: {train_acc*100:.2f}% || "
              f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc*100:.2f}%")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
            epochs_no_improve = 0
        else:
            epochs_no_improve += 1
            if epochs_no_improve >= PATIENCE:
                print(f"Early stopping at epoch {epoch+1} (no improvement for {PATIENCE} epochs)")
                break

    # Restore best weights
    if best_state is not None:
        model.load_state_dict(best_state)

    # Save best model weights
    os.makedirs("models", exist_ok=True)
    model_save_path = os.path.join("models", "solar_panel_resnet50_best.pt")
    torch.save(model.state_dict(), model_save_path)
    print(f"\nBest model saved to {model_save_path}")

    # Generate final classification report
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(DEVICE)
            outputs = model(images)
            _, predicted = outputs.max(1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.numpy())

    print("\n--- Final Validation Classification Report ---")
    print(classification_report(all_labels, all_preds, target_names=classes))

if __name__ == "__main__":
    main()