train.py

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from models.unet import UNet
from data.cityscapes_dataset import CityscapesDataset
from utils import calculate_iou, FocalLoss, dice_loss, boundary_loss, tversky_loss
from config import *
import os
import numpy as np
import matplotlib.pyplot as plt
import random

# Set seeds for reproducibility
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)
np.random.seed(42)
random.seed(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

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

# Data augmentations for training - Conservative for better stability
train_input_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.02),  # Very conservative
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

train_mask_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE, interpolation=transforms.InterpolationMode.NEAREST),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.PILToTensor()
])

# Validation transforms
val_input_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

val_mask_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE, interpolation=transforms.InterpolationMode.NEAREST),
    transforms.PILToTensor()
])

# Datasets and Loaders
train_dataset = CityscapesDataset(DATASET_PATH, "train", train_input_transform, train_mask_transform)
val_dataset = CityscapesDataset(DATASET_PATH, "val", val_input_transform, val_mask_transform)

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

# Model
model = UNet(NUM_CLASSES).to(device)

# Loss and Optimizer - Using combined loss for better segmentation
focal_loss = FocalLoss(alpha=1, gamma=2, ignore_index=255)
ce_loss = nn.CrossEntropyLoss(ignore_index=255)

def combined_loss(pred, target):
    """Simplified combined loss for better convergence"""
    ce = ce_loss(pred, target)
    focal = focal_loss(pred, target)
    dice = dice_loss(pred, target)
    
    # Simplified combination - less conflicting gradients
    return 0.5 * ce + 0.3 * focal + 0.2 * dice

optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4, betas=(0.9, 0.999))
# Use ReduceLROnPlateau for better convergence
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=5, verbose=True, min_lr=1e-7)

# Training Loop
train_losses = []
val_losses = []
val_ious = []
best_val_iou = 0.0
patience_counter = 0
patience = 10  # Early stopping patience

print("Training started...")
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"Training samples: {len(train_dataset)}")
print(f"Validation samples: {len(val_dataset)}")

# Gradient accumulation for effective larger batch size
accumulation_steps = 2  # Effective batch size = BATCH_SIZE * accumulation_steps

for epoch in range(EPOCHS):
    model.train()
    running_loss = 0.0
    optimizer.zero_grad()
    
    for batch_idx, (images, masks) in enumerate(train_loader):
        images = images.to(device)
        masks = masks.squeeze(1).long().to(device)

        outputs = model(images)
        loss = combined_loss(outputs, masks) / accumulation_steps  # Scale loss

        loss.backward()
        
        if (batch_idx + 1) % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()

        running_loss += loss.item() * accumulation_steps  # Unscale for logging

    avg_train_loss = running_loss / len(train_loader)
    train_losses.append(avg_train_loss)

    # Validation with Test-Time Augmentation
    model.eval()
    val_loss = 0.0
    total_iou = 0
    with torch.no_grad():
        for images, masks in val_loader:
            images = images.to(device)
            masks = masks.squeeze(1).long().to(device)

            outputs = model(images)
            loss = ce_loss(outputs, masks)  # Use CE loss for validation
            val_loss += loss.item()

            # Test-Time Augmentation for better validation IoU
            images_flip = torch.flip(images, dims=[3])
            outputs_flip = model(images_flip)
            outputs_flip = torch.flip(outputs_flip, dims=[3])
            outputs_tta = (outputs + outputs_flip) / 2

            preds = torch.argmax(outputs_tta, dim=1)
            iou = calculate_iou(preds, masks, NUM_CLASSES, ignore_index=255)
            total_iou += iou

    avg_val_loss = val_loss / len(val_loader)
    avg_val_iou = total_iou / len(val_loader)
    val_losses.append(avg_val_loss)
    val_ious.append(avg_val_iou)

    scheduler.step(avg_val_iou)  # ReduceLROnPlateau needs validation metric

    print(f"Epoch {epoch+1}/{EPOCHS}, Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}, Val IoU: {avg_val_iou:.4f}, LR: {optimizer.param_groups[0]['lr']:.2e}")

    if avg_val_iou > best_val_iou:
        best_val_iou = avg_val_iou
        patience_counter = 0
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict(),
            'best_val_iou': best_val_iou,
            'val_loss': avg_val_loss,
            'train_loss': avg_train_loss,
            'config': {
                'num_classes': NUM_CLASSES,
                'learning_rate': LEARNING_RATE,
                'batch_size': BATCH_SIZE,
                'image_size': IMAGE_SIZE
            }
        }, "unet_agritech_best.pth")
        print(f"Best model saved with IoU: {best_val_iou:.4f}")
    else:
        patience_counter += 1
        
    # Early stopping
    if patience_counter >= patience:
        print(f"Early stopping triggered after {epoch+1} epochs")
        break

# Plot results
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Train Loss')
plt.plot(val_losses, label='Val Loss')
plt.title("Loss Curve")
plt.legend()
plt.grid()

plt.subplot(1, 2, 2)
plt.plot(val_ious, label='Val IoU')
plt.title("Validation IoU")
plt.legend()
plt.grid()

# Save final model
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'final_val_iou': avg_val_iou,
    'best_val_iou': best_val_iou,
    'train_losses': train_losses,
    'val_losses': val_losses,
    'val_ious': val_ious,
    'config': {
        'num_classes': NUM_CLASSES,
        'learning_rate': LEARNING_RATE,
        'batch_size': BATCH_SIZE,
        'image_size': IMAGE_SIZE
    }
}, "unet_agritech.pth")

plt.tight_layout()
plt.savefig("training_curves.png")
print("Training completed. Curves saved as 'training_curves.png'")
print(f"Final model saved as 'unet_agritech.pth'")
print(f"Best validation IoU achieved: {best_val_iou:.4f}")