simclr_model.py

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm

# Import the architecture and data preprocessing modules
from simclr_architecture import UnderwaterAcousticSimCLR, NTXentLoss
from data_preprocessing import create_data_loaders

class SimCLRModel:
    """
    SimCLR model implementation for underwater acoustic spectrograms.
    Combines the architecture with training functionality.
    """
    def __init__(self, config):
        """
        Initialize the SimCLR model.
        
        Args:
            config: Dictionary containing model configuration
        """
        self.config = config
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Using device: {self.device}")
        
        # Create model
        self.model = UnderwaterAcousticSimCLR(
            base_model=config['base_model'],
            pretrained=config['pretrained'],
            projection_dim=config['projection_dim']
        ).to(self.device)
        
        # Create loss function
        self.criterion = NTXentLoss(
            temperature=config['temperature'],
            batch_size=config['batch_size']
        )
        
        # Create optimizer
        self.optimizer = optim.Adam(
            self.model.parameters(),
            lr=config['learning_rate'],
            weight_decay=config['weight_decay']
        )
        
        # Create learning rate scheduler
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer,
            T_max=config['epochs']
        )
        
        # Initialize tracking variables
        self.current_epoch = 0
        self.train_losses = []
        self.val_losses = []
        
    def train(self, train_loader, val_loader, epochs):
        """
        Train the SimCLR model.
        
        Args:
            train_loader: DataLoader for training data
            val_loader: DataLoader for validation data
            epochs: Number of epochs to train for
        """
        print(f"Starting training for {epochs} epochs...")
        
        for epoch in range(epochs):
            self.current_epoch = epoch
            
            # Training phase
            self.model.train()
            train_loss = 0.0
            train_steps = 0
            
            for i, (img_i, img_j) in enumerate(tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")):
                # Move data to device
                img_i = img_i.to(self.device)
                img_j = img_j.to(self.device)
                
                # Forward pass
                _, z_i = self.model(img_i)
                _, z_j = self.model(img_j)
                
                # Compute loss
                loss = self.criterion(z_i, z_j)
                
                # Backward pass and optimize
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
                
                # Update statistics
                train_loss += loss.item()
                train_steps += 1
                
                # Print progress
                if (i + 1) % 10 == 0:
                    print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}")
            
            # Calculate average training loss
            avg_train_loss = train_loss / train_steps
            self.train_losses.append(avg_train_loss)
            
            # Validation phase
            self.model.eval()
            val_loss = 0.0
            val_steps = 0
            
            with torch.no_grad():
                for img_i, img_j in tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} [Val]"):
                    # Move data to device
                    img_i = img_i.to(self.device)
                    img_j = img_j.to(self.device)
                    
                    # Forward pass
                    _, z_i = self.model(img_i)
                    _, z_j = self.model(img_j)
                    
                    # Compute loss
                    loss = self.criterion(z_i, z_j)
                    
                    # Update statistics
                    val_loss += loss.item()
                    val_steps += 1
            
            # Calculate average validation loss
            avg_val_loss = val_loss / val_steps
            self.val_losses.append(avg_val_loss)
            
            # Update learning rate
            self.scheduler.step()
            
            # Print epoch summary
            print(f"Epoch [{epoch+1}/{epochs}], Train Loss: {avg_train_loss:.4f}, Val Loss: {avg_val_loss:.4f}")
            
            # Save checkpoint
            if (epoch + 1) % 10 == 0 or (epoch + 1) == epochs:
                self.save_checkpoint(f"checkpoint_epoch_{epoch+1}.pt")
                
            # Plot and save loss curves
            self.plot_loss_curves()
    
    def save_checkpoint(self, filename):
        """
        Save model checkpoint.
        
        Args:
            filename: Name of the checkpoint file
        """
        checkpoint_path = os.path.join('/home/ubuntu', filename)
        
        checkpoint = {
            'epoch': self.current_epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'train_losses': self.train_losses,
            'val_losses': self.val_losses,
            'config': self.config
        }
        
        torch.save(checkpoint, checkpoint_path)
        print(f"Checkpoint saved to {checkpoint_path}")
    
    def load_checkpoint(self, checkpoint_path):
        """
        Load model checkpoint.
        
        Args:
            checkpoint_path: Path to the checkpoint file
        """
        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        self.current_epoch = checkpoint['epoch']
        self.train_losses = checkpoint['train_losses']
        self.val_losses = checkpoint['val_losses']
        
        print(f"Checkpoint loaded from {checkpoint_path} (epoch {self.current_epoch})")
    
    def plot_loss_curves(self):
        """
        Plot and save training and validation loss curves.
        """
        plt.figure(figsize=(10, 5))
        plt.plot(range(1, len(self.train_losses) + 1), self.train_losses, label='Train Loss')
        plt.plot(range(1, len(self.val_losses) + 1), self.val_losses, label='Validation Loss')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.title('Training and Validation Loss')
        plt.legend()
        plt.grid(True)
        plt.savefig('/home/ubuntu/loss_curves.png')
        plt.close()
    
    def extract_features(self, dataloader):
        """
        Extract features from the backbone network.
        
        Args:
            dataloader: DataLoader for the data
            
        Returns:
            features: Extracted features
            labels: Corresponding labels (if available)
        """
        self.model.eval()
        features = []
        
        with torch.no_grad():
            for batch in tqdm(dataloader, desc="Extracting features"):
                # Handle both SimCLR mode and evaluation mode
                if isinstance(batch, tuple) and len(batch) == 2 and isinstance(batch[0], torch.Tensor):
                    # SimCLR mode - just use the first view
                    img = batch[0].to(self.device)
                else:
                    # Evaluation mode
                    img = batch.to(self.device)
                
                # Extract features from the backbone
                feature, _ = self.model(img)
                features.append(feature.cpu().numpy())
        
        # Concatenate all features
        features = np.concatenate(features, axis=0)
        
        return features
    
    def save_model(self, path):
        """
        Save the trained model.
        
        Args:
            path: Path to save the model
        """
        torch.save(self.model.state_dict(), path)
        print(f"Model saved to {path}")
    
    def load_model(self, path):
        """
        Load a trained model.
        
        Args:
            path: Path to the model file
        """
        self.model.load_state_dict(torch.load(path, map_location=self.device))
        print(f"Model loaded from {path}")

# Example usage
if __name__ == "__main__":
    # Configuration
    config = {
        'base_model': 'resnet18',
        'pretrained': False,
        'projection_dim': 128,
        'batch_size': 32,
        'temperature': 0.5,
        'learning_rate': 0.0003,
        'weight_decay': 1e-4,
        'epochs': 100,
        'augmentation_probability': 0.5
    }
    
    # Create data loaders
    data_dir = '/home/ubuntu/data'
    train_loader, val_loader = create_data_loaders(
        data_dir=data_dir,
        batch_size=config['batch_size'],
        num_workers=4,
        simclr_mode=True
    )
    
    # Create and train model
    model = SimCLRModel(config)
    
    # Print model summary
    print("Model architecture:")
    print(model.model)
    
    print("SimCLR model implementation complete. Ready for training.")