image-classification-tools 0.5.6__py3-none-any.whl

image_classification_tools/pytorch/hyperparameter_optimization.py

@@ -0,0 +1,315 @@
+'''Hyperparameter optimization utilities for CNN models using Optuna.
+
+This module provides functions for building configurable CNN architectures
+and running hyperparameter optimization with Optuna.
+'''
+
+from typing import Callable
+
+import optuna
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torchvision import datasets
+
+from image_classification_tools.pytorch.data import (
+    load_dataset, prepare_splits, create_dataloaders
+)
+
+
+def create_cnn(
+    n_conv_blocks: int,
+    initial_filters: int,
+    n_fc_layers: int,
+    conv_dropout_rate: float,
+    fc_dropout_rate: float,
+    num_classes: int,
+    in_channels: int = 3
+) -> nn.Sequential:
+    '''Create a CNN with configurable architecture.
+    
+    This function builds a flexible CNN architecture with conv blocks that
+    progressively double filters (initial_filters -> 2x -> 4x -> 8x, etc.).
+    Each conv block contains: 2 Conv layers + BatchNorm + ReLU + MaxPool + Dropout.
+    Uses adaptive pooling before classifier to handle variable spatial dimensions.
+    
+    Args:
+        n_conv_blocks: Number of convolutional blocks (1-5)
+        initial_filters: Number of filters in first conv block (doubles each block)
+        n_fc_layers: Number of fully connected layers before output (1-4)
+        conv_dropout_rate: Dropout probability after convolutional blocks
+        fc_dropout_rate: Dropout probability in fully connected layers
+        num_classes: Number of output classes (required)
+        in_channels: Number of input channels (default: 3 for RGB images)
+    
+    Returns:
+        nn.Sequential model
+    '''
+
+    layers = []
+    current_channels = in_channels
+    
+    # Convolutional blocks
+    for block_idx in range(n_conv_blocks):
+        out_channels = initial_filters * (2 ** block_idx)
+        
+        # First conv in block
+        layers.append(nn.Conv2d(current_channels, out_channels, kernel_size=3, padding=1))
+        layers.append(nn.BatchNorm2d(out_channels))
+        layers.append(nn.ReLU())
+        
+        # Second conv in block
+        layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))
+        layers.append(nn.BatchNorm2d(out_channels))
+        layers.append(nn.ReLU())
+        
+        # Pooling and dropout
+        layers.append(nn.MaxPool2d(2, 2))
+        layers.append(nn.Dropout(conv_dropout_rate))
+        
+        current_channels = out_channels
+    
+    # Classifier with adaptive pooling
+    layers.append(nn.AdaptiveAvgPool2d((1, 1)))
+    layers.append(nn.Flatten())
+    
+    # Generate FC layer sizes by halving from current_channels
+    fc_sizes = []
+    current_fc_size = current_channels // 2
+    for _ in range(n_fc_layers):
+        fc_sizes.append(max(32, current_fc_size))  # Minimum 32 units
+        current_fc_size //= 2
+    
+    # Add FC layers
+    in_features = current_channels
+    for fc_size in fc_sizes:
+        layers.append(nn.Linear(in_features, fc_size))
+        layers.append(nn.ReLU())
+        layers.append(nn.Dropout(fc_dropout_rate))
+        in_features = fc_size
+    
+    # Output layer
+    layers.append(nn.Linear(in_features, num_classes))
+    
+    return nn.Sequential(*layers)
+
+
+def train_trial(
+    model: nn.Module,
+    optimizer: optim.Optimizer,
+    criterion: nn.Module,
+    train_loader: DataLoader,
+    val_loader: DataLoader,
+    n_epochs: int,
+    trial: optuna.Trial
+) -> float:
+    '''Train a model for a single Optuna trial with pruning support.
+    
+    Args:
+        model: PyTorch model to train
+        optimizer: Optimizer for training
+        criterion: Loss function
+        train_loader: DataLoader for training data
+        val_loader: DataLoader for validation data
+        n_epochs: Number of epochs to train
+        trial: Optuna trial object for reporting and pruning
+    
+    Returns:
+        Best validation accuracy achieved during training
+    '''
+    best_val_accuracy = 0.0
+    
+    for epoch in range(n_epochs):
+
+        # Training phase
+        model.train()
+
+        for images, labels in train_loader:
+            optimizer.zero_grad()
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+        
+        # Validation phase
+        model.eval()
+        val_correct = 0
+        val_total = 0
+        
+        with torch.no_grad():
+            for images, labels in val_loader:
+                outputs = model(images)
+                _, predicted = torch.max(outputs.data, 1)
+                val_total += labels.size(0)
+                val_correct += (predicted == labels).sum().item()
+        
+        val_accuracy = 100 * val_correct / val_total
+        best_val_accuracy = max(best_val_accuracy, val_accuracy)
+        
+        # Report intermediate value for pruning
+        trial.report(val_accuracy, epoch)
+        
+        # Prune unpromising trials
+        if trial.should_prune():
+            raise optuna.TrialPruned()
+    
+    return best_val_accuracy
+
+
+def create_objective(
+    data_dir,
+    transform,
+    n_epochs: int,
+    device: torch.device,
+    num_classes: int,
+    in_channels: int = 3,
+    search_space: dict = None
+) -> Callable[[optuna.Trial], float]:
+    '''Create an Optuna objective function for CNN hyperparameter optimization.
+    
+    This factory function creates a closure that captures the data loading parameters
+    and training configuration, returning an objective function suitable for Optuna.
+    
+    Args:
+        data_dir: Directory containing training data
+        transform: Transform to apply to both training and validation data
+        n_epochs: Number of epochs per trial
+        device: Device to train on (cuda or cpu)
+        num_classes: Number of output classes (required, e.g., 10 for CIFAR-10)
+        in_channels: Number of input channels (default: 3 for RGB images, 1 for grayscale)
+        search_space: Dictionary defining hyperparameter search space (default: None)
+    
+    Returns:
+        Objective function for optuna.Study.optimize()
+    
+    Example:
+        >>> objective = create_objective(
+        ...     data_dir='data/', 
+        ...     transform=transform,
+        ...     n_epochs=50, 
+        ...     device=device,
+        ...     num_classes=10
+        ... )
+        >>> study = optuna.create_study(direction='maximize')
+        >>> study.optimize(objective, n_trials=100)
+    '''
+    
+    # Default search space if none provided
+    if search_space is None:
+        search_space = {
+            'batch_size': [64, 128, 256, 512],
+            'n_conv_blocks': (1, 5),
+            'initial_filters': [16, 32, 64, 128],
+            'n_fc_layers': (1, 4),
+            'conv_dropout_rate': (0.1, 0.5),
+            'fc_dropout_rate': (0.3, 0.7),
+            'learning_rate': (1e-5, 1e-2, 'log'),
+            'optimizer': ['Adam', 'SGD', 'RMSprop'],
+            'sgd_momentum': (0.8, 0.99),
+            'weight_decay': (1e-6, 1e-3, 'log')
+        }
+    
+    def objective(trial: optuna.Trial) -> float:
+        '''Optuna objective function for CNN hyperparameter optimization.'''
+        
+        # Suggest hyperparameters from search space
+        batch_size = trial.suggest_categorical('batch_size', search_space['batch_size'])
+        n_conv_blocks = trial.suggest_int('n_conv_blocks', *search_space['n_conv_blocks'])
+        initial_filters = trial.suggest_categorical('initial_filters', search_space['initial_filters'])
+        n_fc_layers = trial.suggest_int('n_fc_layers', *search_space['n_fc_layers'])
+        conv_dropout_rate = trial.suggest_float('conv_dropout_rate', *search_space['conv_dropout_rate'])
+        fc_dropout_rate = trial.suggest_float('fc_dropout_rate', *search_space['fc_dropout_rate'])
+        
+        # Handle learning rate with optional log scale
+        lr_params = search_space['learning_rate']
+        learning_rate = trial.suggest_float('learning_rate', lr_params[0], lr_params[1], 
+                                           log=(lr_params[2] == 'log' if len(lr_params) > 2 else False))
+        
+        optimizer_name = trial.suggest_categorical('optimizer', search_space['optimizer'])
+        
+        # Weight decay
+        wd_params = search_space['weight_decay']
+        weight_decay = trial.suggest_float('weight_decay', wd_params[0], wd_params[1],
+                                          log=(wd_params[2] == 'log' if len(wd_params) > 2 else False))
+        
+        # Create data loaders with suggested batch size
+        # Load datasets
+        train_dataset = load_dataset(
+            data_source=datasets.CIFAR10,
+            transform=transform,
+            train=True,
+            download=False,
+            root=data_dir
+        )
+        
+        test_dataset = load_dataset(
+            data_source=datasets.CIFAR10,
+            transform=transform,
+            train=False,
+            download=False,
+            root=data_dir
+        )
+        
+        # Prepare splits
+        train_dataset, val_dataset, _ = prepare_splits(
+            train_dataset=train_dataset,
+            test_dataset=test_dataset,
+            val_size=10000
+        )
+        
+        # Create dataloaders with memory preloading
+        train_loader, val_loader, _ = create_dataloaders(
+            train_dataset, val_dataset, val_dataset,
+            batch_size=batch_size,
+            preload_to_memory=(device is not None),
+            device=device
+        )
+        
+        # Create model with suggested architecture
+        model = create_cnn(
+            n_conv_blocks=n_conv_blocks,
+            initial_filters=initial_filters,
+            n_fc_layers=n_fc_layers,
+            conv_dropout_rate=conv_dropout_rate,
+            fc_dropout_rate=fc_dropout_rate,
+            num_classes=num_classes,
+            in_channels=in_channels
+        ).to(device)
+        
+        # Define optimizer
+        if optimizer_name == 'Adam':
+            optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+
+        elif optimizer_name == 'SGD':
+            momentum = trial.suggest_float('sgd_momentum', *search_space['sgd_momentum'])
+            optimizer = optim.SGD(model.parameters(), lr=learning_rate, 
+                                momentum=momentum, weight_decay=weight_decay)
+        
+        else:  # RMSprop
+            optimizer = optim.RMSprop(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+        
+        criterion = nn.CrossEntropyLoss()
+        
+        # Train model and return best validation accuracy
+        try:
+            return train_trial(
+                model=model,
+                optimizer=optimizer,
+                criterion=criterion,
+                train_loader=train_loader,
+                val_loader=val_loader,
+                n_epochs=n_epochs,
+                trial=trial
+            )
+
+        except RuntimeError as e:
+            # Catch architecture errors (e.g., dimension mismatches)
+            raise optuna.TrialPruned(f'RuntimeError with params: {trial.params} - {str(e)}')
+        
+        except torch.cuda.OutOfMemoryError:
+            # Clear CUDA cache and skip this trial
+            torch.cuda.empty_cache()
+            raise optuna.TrialPruned(f'CUDA OOM with params: {trial.params}')
+    
+    return objective

image_classification_tools/pytorch/plotting.py

@@ -0,0 +1,311 @@
+'''Plotting functions for image classification models.
+
+This module provides visualization utilities for analyzing image classification models,
+including sample images, learning curves, confusion matrices, and evaluation metrics.
+'''
+
+import matplotlib.pyplot as plt
+import numpy as np
+from torch.utils.data import Dataset
+
+
+def plot_sample_images(
+    dataset: Dataset,
+    class_names: list[str],
+    nrows: int = 2,
+    ncols: int = 5,
+    figsize: tuple[float, float] | None = None
+) -> tuple[plt.Figure, np.ndarray]:
+    '''Plot sample images from a dataset.
+    
+    Automatically handles both grayscale (1 channel) and RGB (3 channel) images.
+    
+    Args:
+        dataset: PyTorch dataset containing (image, label) tuples.
+        class_names: List of class names for labeling.
+        nrows: Number of rows in the grid.
+        ncols: Number of columns in the grid.
+        figsize: Figure size (width, height). Defaults to (ncols*1.5, nrows*1.5).
+        
+    Returns:
+        Tuple of (figure, axes array).
+    '''
+    if figsize is None:
+        figsize = (ncols * 1.5, nrows * 1.5)
+    
+    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize)
+    axes = axes.flatten()
+
+    for i, ax in enumerate(axes):
+        # Get image and label from dataset
+        img, label = dataset[i]
+        
+        # Unnormalize for plotting
+        img = img * 0.5 + 0.5
+        img = img.numpy()
+        
+        # Handle grayscale vs RGB images
+        if img.shape[0] == 1:
+
+            # Grayscale: squeeze channel dimension
+            img = img.squeeze()
+            ax.imshow(img, cmap='gray')
+    
+        else:
+
+            # RGB: transpose from (C, H, W) to (H, W, C)
+            img = np.transpose(img, (1, 2, 0))
+            ax.imshow(img)
+        
+        ax.set_title(class_names[label])
+        ax.axis('off')
+
+    plt.tight_layout()
+    
+    return fig, axes
+
+
+def plot_learning_curves(
+    history: dict[str, list[float]],
+    figsize: tuple[float, float] = (10, 4)
+) -> tuple[plt.Figure, np.ndarray]:
+    '''Plot training and validation loss and accuracy curves.
+    
+    Args:
+        history: Dictionary containing 'train_loss', 'val_loss', 
+                 'train_accuracy', and 'val_accuracy' lists.
+        figsize: Figure size (width, height).
+        
+    Returns:
+        Tuple of (figure, axes array).
+    '''
+    fig, axes = plt.subplots(1, 2, figsize=figsize)
+
+    axes[0].set_title('Loss')
+    axes[0].plot(history['train_loss'], label='Train')
+    axes[0].plot(history['val_loss'], label='Validation')
+    axes[0].set_xlabel('Epoch')
+    axes[0].set_ylabel('Loss (cross-entropy)')
+    axes[0].legend(loc='best')
+
+    axes[1].set_title('Accuracy')
+    axes[1].plot(history['train_accuracy'], label='Train')
+    axes[1].plot(history['val_accuracy'], label='Validation')
+    axes[1].set_xlabel('Epoch')
+    axes[1].set_ylabel('Accuracy (%)')
+    axes[1].legend(loc='best')
+
+    plt.tight_layout()
+    
+    return fig, axes
+
+
+def plot_confusion_matrix(
+    true_labels: np.ndarray,
+    predictions: np.ndarray,
+    class_names: list[str],
+    figsize: tuple[float, float] = (8, 8),
+    cmap: str = 'Blues'
+) -> tuple[plt.Figure, plt.Axes]:
+    '''Plot a confusion matrix heatmap.
+    
+    Args:
+        true_labels: Array of true class labels.
+        predictions: Array of predicted class labels.
+        class_names: List of class names for labeling.
+        figsize: Figure size (width, height).
+        cmap: Colormap for the heatmap.
+        
+    Returns:
+        Tuple of (figure, axes).
+    '''
+    from sklearn.metrics import confusion_matrix
+    
+    cm = confusion_matrix(true_labels, predictions)
+    
+    fig, ax = plt.subplots(figsize=figsize)
+
+    ax.set_title('Confusion matrix')
+    im = ax.imshow(cm, cmap=cmap)
+
+    # Add labels
+    ax.set_xticks(range(len(class_names)))
+    ax.set_yticks(range(len(class_names)))
+    ax.set_xticklabels(class_names, rotation=45, ha='right')
+    ax.set_yticklabels(class_names)
+    ax.set_xlabel('Predicted label')
+    ax.set_ylabel('True label')
+
+    # Add text annotations
+    for i in range(len(class_names)):
+        for j in range(len(class_names)):
+            color = 'white' if cm[i, j] > cm.max() / 2 else 'black'
+            ax.text(j, i, str(cm[i, j]), ha='center', va='center', color=color)
+
+    plt.tight_layout()
+    
+    return fig, ax
+
+
+def plot_class_probability_distributions(
+    all_probs: np.ndarray,
+    class_names: list[str],
+    nrows: int = 2,
+    ncols: int = 5,
+    figsize: tuple[float, float] = (12, 4),
+    bins: int = 50,
+    color: str = 'black'
+) -> tuple[plt.Figure, np.ndarray]:
+    '''Plot predicted probability distributions for each class.
+    
+    Args:
+        all_probs: Array of shape (n_samples, n_classes) with predicted probabilities.
+        class_names: List of class names for labeling.
+        nrows: Number of rows in the subplot grid.
+        ncols: Number of columns in the subplot grid.
+        figsize: Figure size (width, height).
+        bins: Number of histogram bins.
+        color: Histogram bar color.
+        
+    Returns:
+        Tuple of (figure, axes array).
+    '''
+    fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=figsize)
+
+    fig.suptitle('Predicted probability distributions by class', fontsize=14, y=1.02)
+    fig.supxlabel('Predicted probability', fontsize=12)
+    fig.supylabel('Count', fontsize=12)
+
+    axes = axes.flatten()
+
+    for i, (ax, class_name) in enumerate(zip(axes, class_names)):
+        # Get probabilities for this class across all samples
+        class_probs = all_probs[:, i]
+        
+        # Plot histogram
+        ax.hist(class_probs, bins=bins, color=color)
+        ax.set_title(class_name)
+        ax.set_xlim(0, 1)
+
+    plt.tight_layout()
+    
+    return fig, axes
+
+
+def plot_evaluation_curves(
+    true_labels: np.ndarray,
+    all_probs: np.ndarray,
+    class_names: list[str],
+    figsize: tuple[float, float] = (12, 5)
+) -> tuple[plt.Figure, tuple[plt.Axes, plt.Axes]]:
+    '''Plot ROC and Precision-Recall curves for multi-class classification.
+    
+    Args:
+        true_labels: Array of true class labels.
+        all_probs: Array of shape (n_samples, n_classes) with predicted probabilities.
+        class_names: List of class names for labeling.
+        figsize: Figure size (width, height).
+        
+    Returns:
+        Tuple of (figure, (ax1, ax2)).
+    '''
+    from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
+    from sklearn.preprocessing import label_binarize
+    
+    # Binarize true labels for one-vs-rest evaluation
+    y_test_bin = label_binarize(true_labels, classes=range(len(class_names)))
+
+    # Create figure with ROC and PR curves side by side
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
+
+    # Plot ROC curves for each class
+    ax1.set_title('ROC curves (one-vs-rest)')
+
+    for i, class_name in enumerate(class_names):
+        fpr, tpr, _ = roc_curve(y_test_bin[:, i], all_probs[:, i])
+        roc_auc = auc(fpr, tpr)
+        ax1.plot(fpr, tpr, label=class_name)
+
+    ax1.plot([0, 1], [0, 1], 'k--', label='random classifier')
+    ax1.set_xlabel('False positive rate')
+    ax1.set_ylabel('True positive rate')
+    ax1.legend(loc='lower right', fontsize=12)
+    ax1.set_xlim([0, 1])
+    ax1.set_ylim([0, 1.05])
+
+    # Plot Precision-Recall curves for each class
+    ax2.set_title('Precision-recall curves (one-vs-rest)')
+
+    for i, class_name in enumerate(class_names):
+        precision, recall, _ = precision_recall_curve(y_test_bin[:, i], all_probs[:, i])
+        ap = average_precision_score(y_test_bin[:, i], all_probs[:, i])
+        ax2.plot(recall, precision)
+
+    # Random classifier baseline (horizontal line at class prevalence = 1/num_classes)
+    baseline = 1 / len(class_names)
+    ax2.axhline(y=baseline, color='k', linestyle='--')
+
+    ax2.set_xlabel('Recall')
+    ax2.set_ylabel('Precision')
+    ax2.set_xlim([0, 1])
+    ax2.set_ylim([0, 1.05])
+
+    plt.tight_layout()
+    
+    return fig, (ax1, ax2)
+
+
+def plot_optimization_results(
+    study,
+    figsize: tuple[float, float] = (12, 4)
+) -> tuple[plt.Figure, np.ndarray]:
+    '''Plot Optuna optimization history and hyperparameter importance.
+    
+    Args:
+        study: Optuna study object with completed trials.
+        figsize: Figure size (width, height).
+        
+    Returns:
+        Tuple of (figure, axes array).
+    '''
+    import optuna
+    
+    fig, axes = plt.subplots(1, 2, figsize=figsize)
+
+    # Optimization history
+    axes[0].set_title('Optimization History')
+    
+    trial_numbers = [t.number for t in study.trials if t.value is not None]
+    trial_values = [t.value for t in study.trials if t.value is not None]
+
+    axes[0].plot(trial_numbers, trial_values, 'ko-', alpha=0.6)
+    axes[0].axhline(
+        y=study.best_value,
+        color='r', linestyle='--', label=f'Best: {study.best_value:.2f}%'
+    )
+    axes[0].set_xlabel('Trial')
+    axes[0].set_ylabel('Validation Accuracy (%)')
+    axes[0].legend()
+
+    # Hyperparameter importance (if enough trials completed)
+    axes[1].set_title('Hyperparameter Importance')
+    completed_trials = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
+
+    if len(completed_trials) >= 5:
+        importance = optuna.importance.get_param_importances(study)
+        params = list(importance.keys())
+        values = list(importance.values())
+        
+        axes[1].set_xlabel('Importance')
+        axes[1].barh(params, values, color='black')
+
+    else:
+        axes[1].text(
+            0.5, 0.5,
+            'Not enough completed trials\nfor importance analysis', 
+            ha='center', va='center', transform=axes[1].transAxes
+        )
+
+    plt.tight_layout()
+    
+    return fig, axes