train_oscd_a100.py

# %% [markdown]
# # OSCD Change Detection Training Pipeline
# 
# This script trains a Vision Transformer model for change detection on the OSCD dataset.
# It handles image pairs (before/after) and binary change detection labels.

# %%
import os
import json
import time
import warnings
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Any
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
from tqdm import tqdm

# Local imports
from data_loader_oscd import create_oscd_dataloaders
from models.vision_transformer import create_oscd_model

# %%
class OSCDTrainer:
    """Trainer for OSCD change detection model."""
    
    def __init__(self, config: Dict[str, Any]):
        """Initialize the trainer."""
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Create directories
        self.log_dir = Path(config['log_dir'])
        self.checkpoint_dir = Path(config['checkpoint_dir'])
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        
        # Initialize components
        self.model = None
        self.optimizer = None
        self.scheduler = None
        self.criterion = None
        self.train_loader = None
        self.val_loader = None
        self.test_loader = None
        self.writer = None
        
        print(f"OSCD Trainer initialized on device: {self.device}")
    
    def setup_data(self):
        """Setup data loaders."""
        print("Setting up data loaders...")
        
        self.train_loader, self.val_loader, self.test_loader, metadata = create_oscd_dataloaders(
            data_dir=self.config['data_dir'],
            batch_size=self.config['batch_size'],
            train_split=self.config['train_split'],
            val_split=self.config['val_split'],
            test_split=self.config['test_split'],
            num_workers=self.config['num_workers'],
            use_augmentation=self.config['use_augmentation'],
            seed=self.config['seed'],
            target_size=self.config['target_size']
        )
        
        self.metadata = metadata
        print(f"Data loaded: {len(self.train_loader.dataset)} train, "
              f"{len(self.val_loader.dataset)} val, {len(self.test_loader.dataset)} test samples")
    
    def setup_model(self):
        """Setup model, optimizer, and loss function."""
        print("Setting up model...")
        
        # Create model
        self.model = create_oscd_model(
            model_size=self.config['model_size'],
            use_pretrained=self.config['use_pretrained']
        ).to(self.device)
        
        # Loss function for binary change detection
        self.criterion = nn.BCEWithLogitsLoss()
        
        # Optimizer
        if self.config['optimizer'] == 'adamw':
            self.optimizer = optim.AdamW(
                self.model.parameters(),
                lr=self.config['learning_rate'],
                weight_decay=self.config['weight_decay']
            )
        else:
            self.optimizer = optim.Adam(
                self.model.parameters(),
                lr=self.config['learning_rate'],
                weight_decay=self.config['weight_decay']
            )
        
        # Scheduler
        if self.config['scheduler'] == 'cosine':
            self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer,
                T_max=self.config['epochs']
            )
        elif self.config['scheduler'] == 'step':
            self.scheduler = optim.lr_scheduler.StepLR(
                self.optimizer,
                step_size=30,
                gamma=0.1
            )
        
        print(f"Model created: {self.config['model_size']} size")
        print(f"Optimizer: {self.config['optimizer']}")
        print(f"Scheduler: {self.config['scheduler']}")
    
    def setup_logging(self):
        """Setup TensorBoard logging."""
        self.writer = SummaryWriter(self.log_dir)
        
        # Log configuration
        config_str = json.dumps(self.config, indent=2)
        self.writer.add_text('Configuration', config_str, 0)
    
    def train_epoch(self, epoch: int) -> Dict[str, float]:
        """Train for one epoch."""
        self.model.train()
        total_loss = 0.0
        all_predictions = []
        all_labels = []
        
        pbar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}")
        
        for batch_idx, (image_pairs, labels) in enumerate(pbar):
            # Handle variable-sized images - process one by one
            batch_loss = 0.0
            batch_predictions = []
            batch_labels = []
            
            for i, (image_pair, binary_label) in enumerate(zip(image_pairs, labels)):
                # Ensure binary_label is a scalar tensor
                if binary_label.numel() > 1:
                    binary_label = binary_label.mean()  # Take mean if multiple elements
                binary_label = binary_label.view(1)  # Ensure it's a 1D tensor
                # Ensure binary_label is a scalar tensor
                if binary_label.numel() > 1:
                    binary_label = binary_label.mean()  # Take mean if multiple elements
                binary_label = binary_label.view(1)  # Ensure it's a 1D tensor
                # Move data to device
                image_pair = image_pair.to(self.device)  # Shape: (2, 13, H, W)
                binary_label = binary_label.to(self.device)  # Shape: (1,) - already processed
                
                # Debug: Print shapes for first few samples
                if batch_idx == 0 and i == 0:
                    print(f"DEBUG - Sample {i}:")
                    print(f"  image_pair shape: {image_pair.shape}")
                    print(f"  binary_label shape: {binary_label.shape}")
                    print(f"  binary_label value: {binary_label.item()}")
                
                # Add batch dimension
                image_pair = image_pair.unsqueeze(0)  # Shape: (1, 2, 13, H, W)
                
                # Forward pass
                self.optimizer.zero_grad()
                output = self.model(image_pair)  # Shape: (1, 1) - single scalar per sample
                
                # Debug: Print output shape
                if batch_idx == 0 and i == 0:
                    print(f"  output shape: {output.shape}")
                    print(f"  output.squeeze() shape: {output.squeeze().shape}")
                
                # Calculate loss - ensure both tensors have the same shape
                output_squeezed = output.squeeze()  # Shape: (1,) or scalar
                if output_squeezed.ndim == 0:  # If scalar, make it 1D
                    output_squeezed = output_squeezed.unsqueeze(0)  # Shape: (1,)
                elif output_squeezed.ndim > 1:  # If more than 1D, flatten
                    output_squeezed = output_squeezed.flatten()  # Shape: (N,)
                
                # Ensure binary_label is also 1D
                if binary_label.ndim == 0:  # If scalar, make it 1D
                    binary_label = binary_label.unsqueeze(0)  # Shape: (1,)
                
                # Ensure both have the same number of elements
                if output_squeezed.numel() != binary_label.numel():
                    # If output has multiple elements, take the first one
                    if output_squeezed.numel() > 1:
                        output_squeezed = output_squeezed[0].unsqueeze(0)
                    # If binary_label has multiple elements, take the first one
                    elif binary_label.numel() > 1:
                        binary_label = binary_label[0].unsqueeze(0)
                
                # Debug: Print final shapes
                if batch_idx == 0 and i == 0:
                    print(f"  output_squeezed shape: {output_squeezed.shape}")
                    print(f"  binary_label shape: {binary_label.shape}")
                    print(f"  Shapes match: {output_squeezed.shape == binary_label.shape}")
                    print(f"  Elements match: {output_squeezed.numel() == binary_label.numel()}")
                
                loss = self.criterion(output_squeezed, binary_label)
                
                # Get prediction
                prediction = (torch.sigmoid(output_squeezed) > 0.5).float()
                
                # Backward pass
                loss.backward()
                
                # Gradient clipping
                if self.config['gradient_clip'] > 0:
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['gradient_clip'])
                
                self.optimizer.step()
                
                # Collect metrics
                batch_loss += loss.item()
                batch_predictions.append(prediction.cpu().numpy())
                batch_labels.append(binary_label.cpu().numpy())
            
            # Average loss for the batch
            avg_loss = batch_loss / len(image_pairs)
            total_loss += avg_loss
            
            # Extend predictions and labels
            all_predictions.extend(batch_predictions)
            all_labels.extend(batch_labels)
            
            # Update progress bar
            pbar.set_postfix({'loss': f'{avg_loss:.4f}'})
        
        # Calculate metrics
        accuracy = accuracy_score(all_labels, all_predictions)
        f1 = f1_score(all_labels, all_predictions, average='binary')
        precision = precision_score(all_labels, all_predictions, average='binary', zero_division=0)
        recall = recall_score(all_labels, all_predictions, average='binary', zero_division=0)
        
        metrics = {
            'loss': total_loss / len(self.train_loader),
            'accuracy': accuracy,
            'f1': f1,
            'precision': precision,
            'recall': recall
        }
        
        return metrics
    
    def validate_epoch(self, epoch: int) -> Dict[str, float]:
        """Validate for one epoch."""
        self.model.eval()
        total_loss = 0.0
        all_predictions = []
        all_labels = []
        
        with torch.no_grad():
            for image_pairs, labels in tqdm(self.val_loader, desc=f"Validation {epoch+1}"):
                # Handle variable-sized images - process one by one
                batch_loss = 0.0
                batch_predictions = []
                batch_labels = []
                
                for image_pair, binary_label in zip(image_pairs, labels):
                    # Ensure binary_label is a scalar tensor
                    if binary_label.numel() > 1:
                        binary_label = binary_label.mean()  # Take mean if multiple elements
                    binary_label = binary_label.view(1)  # Ensure it's a 1D tensor
                    
                    # Move data to device
                    image_pair = image_pair.to(self.device)  # Shape: (2, 13, H, W)
                    binary_label = binary_label.to(self.device)  # Shape: (1,) - already processed
                    
                    # Add batch dimension
                    image_pair = image_pair.unsqueeze(0)  # Shape: (1, 2, 13, H, W)
                    
                    output = self.model(image_pair)  # Shape: (1, 1) - single scalar per sample
                    
                    # Calculate loss - ensure both tensors have the same shape
                    output_squeezed = output.squeeze()  # Shape: (1,) or scalar
                    if output_squeezed.ndim == 0:  # If scalar, make it 1D
                        output_squeezed = output_squeezed.unsqueeze(0)  # Shape: (1,)
                    elif output_squeezed.ndim > 1:  # If more than 1D, flatten
                        output_squeezed = output_squeezed.flatten()  # Shape: (N,)
                    
                    # Ensure binary_label is also 1D
                    if binary_label.ndim == 0:  # If scalar, make it 1D
                        binary_label = binary_label.unsqueeze(0)  # Shape: (1,)
                    
                    # Ensure both have the same number of elements
                    if output_squeezed.numel() != binary_label.numel():
                        # If output has multiple elements, take the first one
                        if output_squeezed.numel() > 1:
                            output_squeezed = output_squeezed[0].unsqueeze(0)
                        # If binary_label has multiple elements, take the first one
                        elif binary_label.numel() > 1:
                            binary_label = binary_label[0].unsqueeze(0)
                    
                    loss = self.criterion(output_squeezed, binary_label)
                    
                    # Get prediction
                    prediction = (torch.sigmoid(output_squeezed) > 0.5).float()
                    
                    batch_loss += loss.item()
                    batch_predictions.append(prediction.cpu().numpy())
                    batch_labels.append(binary_label.cpu().numpy())
                
                # Average loss for the batch
                avg_loss = batch_loss / len(image_pairs)
                total_loss += avg_loss
                
                # Extend predictions and labels
                all_predictions.extend(batch_predictions)
                all_labels.extend(batch_labels)
        
        # Calculate metrics
        accuracy = accuracy_score(all_labels, all_predictions)
        f1 = f1_score(all_labels, all_predictions, average='binary')
        precision = precision_score(all_labels, all_predictions, average='binary', zero_division=0)
        recall = recall_score(all_labels, all_predictions, average='binary', zero_division=0)
        
        metrics = {
            'loss': total_loss / len(self.val_loader),
            'accuracy': accuracy,
            'f1': f1,
            'precision': precision,
            'recall': recall
        }
        
        return metrics
    
    def train(self):
        """Main training loop."""
        print("Starting training...")
        
        best_f1 = 0.0
        patience_counter = 0
        
        for epoch in range(self.config['epochs']):
            # Train
            train_metrics = self.train_epoch(epoch)
            
            # Validate
            val_metrics = self.validate_epoch(epoch)
            
            # Update scheduler
            if self.scheduler:
                self.scheduler.step()
            
            # Log metrics
            if self.writer:
                for key, value in train_metrics.items():
                    self.writer.add_scalar(f'Train/{key}', value, epoch)
                for key, value in val_metrics.items():
                    self.writer.add_scalar(f'Val/{key}', value, epoch)
                self.writer.add_scalar('LR', self.optimizer.param_groups[0]['lr'], epoch)
            
            # Print progress
            print(f"Epoch {epoch+1}/{self.config['epochs']}: "
                  f"Train Loss: {train_metrics['loss']:.4f}, "
                  f"Val Loss: {val_metrics['loss']:.4f}, "
                  f"Val F1: {val_metrics['f1']:.4f}")
            
            # Save best model
            if val_metrics['f1'] > best_f1:
                best_f1 = val_metrics['f1']
                patience_counter = 0
                
                checkpoint = {
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,
                    'best_f1': best_f1,
                    'config': self.config
                }
                
                torch.save(checkpoint, self.checkpoint_dir / 'best_checkpoint.pth')
                print(f"New best model saved! F1: {best_f1:.4f}")
            else:
                patience_counter += 1
            
            # Early stopping
            if patience_counter >= self.config['patience']:
                print(f"Early stopping after {epoch+1} epochs")
                break
        
        print(f"Training completed! Best F1: {best_f1:.4f}")
    
    def evaluate(self) -> Dict[str, Any]:
        """Evaluate the model on test set."""
        print("Evaluating on test set...")
        
        # Load best model
        checkpoint_path = self.checkpoint_dir / 'best_checkpoint.pth'
        if checkpoint_path.exists():
            checkpoint = torch.load(checkpoint_path, map_location=self.device)
            self.model.load_state_dict(checkpoint['model_state_dict'])
            print(f"Loaded best model from epoch {checkpoint['epoch']}")
        
        self.model.eval()
        total_loss = 0.0
        all_predictions = []
        all_labels = []
        
        with torch.no_grad():
            for image_pairs, labels in tqdm(self.test_loader, desc="Testing"):
                # Handle variable-sized images - process one by one
                batch_loss = 0.0
                batch_predictions = []
                batch_labels = []
                
                for image_pair, binary_label in zip(image_pairs, labels):
                    # Ensure binary_label is a scalar tensor
                    if binary_label.numel() > 1:
                        binary_label = binary_label.mean()  # Take mean if multiple elements
                    binary_label = binary_label.view(1)  # Ensure it's a 1D tensor
                    
                    # Move data to device
                    image_pair = image_pair.to(self.device)  # Shape: (2, 13, H, W)
                    binary_label = binary_label.to(self.device)  # Shape: (1,) - already processed
                    
                    # Add batch dimension
                    image_pair = image_pair.unsqueeze(0)  # Shape: (1, 2, 13, H, W)
                    
                    output = self.model(image_pair)  # Shape: (1, 1) - single scalar per sample
                    
                    # Calculate loss - ensure both tensors have the same shape
                    output_squeezed = output.squeeze()  # Shape: (1,) or scalar
                    if output_squeezed.ndim == 0:  # If scalar, make it 1D
                        output_squeezed = output_squeezed.unsqueeze(0)  # Shape: (1,)
                    elif output_squeezed.ndim > 1:  # If more than 1D, flatten
                        output_squeezed = output_squeezed.flatten()  # Shape: (N,)
                    
                    # Ensure binary_label is also 1D
                    if binary_label.ndim == 0:  # If scalar, make it 1D
                        binary_label = binary_label.unsqueeze(0)  # Shape: (1,)
                    
                    # Ensure both have the same number of elements
                    if output_squeezed.numel() != binary_label.numel():
                        # If output has multiple elements, take the first one
                        if output_squeezed.numel() > 1:
                            output_squeezed = output_squeezed[0].unsqueeze(0)
                        # If binary_label has multiple elements, take the first one
                        elif binary_label.numel() > 1:
                            binary_label = binary_label[0].unsqueeze(0)
                    
                    loss = self.criterion(output_squeezed, binary_label)
                    
                    # Get prediction
                    prediction = (torch.sigmoid(output_squeezed) > 0.5).float()
                    
                    batch_loss += loss.item()
                    batch_predictions.append(prediction.cpu().numpy())
                    batch_labels.append(binary_label.cpu().numpy())
                
                # Average loss for the batch
                avg_loss = batch_loss / len(image_pairs)
                total_loss += avg_loss
                
                # Extend predictions and labels
                all_predictions.extend(batch_predictions)
                all_labels.extend(batch_labels)
        
        # Calculate metrics
        accuracy = accuracy_score(all_labels, all_predictions)
        f1 = f1_score(all_labels, all_predictions, average='binary')
        precision = precision_score(all_labels, all_predictions, average='binary', zero_division=0)
        recall = recall_score(all_labels, all_predictions, average='binary', zero_division=0)
        
        results = {
            'accuracy': accuracy,
            'f1': f1,
            'precision': precision,
            'recall': recall,
            'loss': total_loss / len(self.test_loader)
        }
        
        print("Test Results:")
        for key, value in results.items():
            print(f"  {key}: {value:.4f}")
        
        return results

# %%
def main():
    """Main training function."""
    
    # Configuration
    config = {
        'model_size': 'base',
        'use_pretrained': False,
        'patch_size': 8,
        'epochs': 100,
        'batch_size': 16,  # Smaller batch size for more stable training
        'learning_rate': 0.0001,  # Lower learning rate for better convergence
        'weight_decay': 0.01,
        'optimizer': 'adamw',
        'scheduler': 'cosine',
        'gradient_clip': 1.0,
        'patience': 20,
        'data_dir': './oscd_npz',
        'train_split': 0.7,
        'val_split': 0.15,
        'test_split': 0.15,
        'use_augmentation': True,
        'num_workers': 8,
        'log_dir': 'logs/oscd_a100',
        'checkpoint_dir': 'checkpoints/oscd_a100',
        'seed': 42,
        'target_size': 64  # Target image size for resizing
        # Removed max_samples_per_city to use all available data
    }
    
    print("=" * 60)
    print("OSCD CHANGE DETECTION TRAINING")
    print("=" * 60)
    
    # Initialize trainer
    trainer = OSCDTrainer(config)
    
    # Setup components
    trainer.setup_data()
    trainer.setup_model()
    trainer.setup_logging()
    
    # Train
    start_time = time.time()
    trainer.train()
    training_time = time.time() - start_time
    
    # Evaluate
    test_results = trainer.evaluate()
    
    # Save results
    results = {
        **test_results,
        'training_time_hours': training_time / 3600,
        'config': config
    }
    
    results_dir = Path('results/oscd_train_results_a100')
    results_dir.mkdir(parents=True, exist_ok=True)
    
    with open(results_dir / 'test_results.json', 'w') as f:
        json.dump(results, f, indent=2, default=str)
    
    print(f"\nResults saved to: {results_dir}")
    print(f"Training time: {training_time/3600:.2f} hours")

# %%
if __name__ == "__main__":
    main()