cifar10-tools 0.2.0__tar.gz → 0.4.0__tar.gz

cifar10_tools-0.4.0/PKG-INFO

@@ -0,0 +1,72 @@
+Metadata-Version: 2.4
+Name: cifar10_tools
+Version: 0.4.0
+Summary: Tools for training neural networks on the CIFAR-10 task with PyTorch and TensorFlow
+License: GPLv3
+License-File: LICENSE
+Keywords: Python,Machine learning,Deep learning,CNNs,Computer vision,Image classification,CIFAR-10
+Author: gperdrizet
+Author-email: george@perdrizet.org
+Requires-Python: >=3.10,<3.13
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
+Classifier: License :: Other/Proprietary License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Provides-Extra: tensorflow
+Requires-Dist: numpy (>=1.24)
+Requires-Dist: torch (>=2.0)
+Requires-Dist: torchvision (>=0.15)
+Project-URL: Documentation, https://gperdrizet.github.io/CIFAR10/README.md
+Project-URL: Homepage, https://github.com/gperdrizet/CIFAR10
+Project-URL: Issues, https://github.com/gperdrizet/CIFAR10/issues
+Project-URL: PyPI, https://pypi.org/project/cifar10_tools
+Project-URL: Repository, https://github.com/gperdrizet/CIFAR10
+Description-Content-Type: text/markdown
+
+# PyTorch: CIFAR-10 Demonstration
+
+A progressive deep learning tutorial for image classification on the CIFAR-10 dataset using PyTorch. This project demonstrates the evolution from basic deep neural networks to optimized convolutional neural networks with data augmentation. It also provides a set of utility functions as a PyPI package for use in other projects.
+
+[View on PyPI](https://pypi.org/project/cifar10_tools)
+
+## Installation
+
+Install the helper tools package locally in editable mode:
+
+```bash
+pip install -e .
+```
+
+## Project Overview
+
+This repository contains a series of Jupyter notebooks that progressively build more sophisticated neural network architectures for the CIFAR-10 image classification task. Each notebook builds upon concepts from the previous one, demonstrating key deep learning techniques.
+
+## Notebooks
+
+| Notebook | Description |
+|----------|-------------|
+| [01-DNN.ipynb](notebooks/01-DNN.ipynb) | **Deep Neural Network** - Baseline fully-connected DNN classifier using `nn.Sequential`. Establishes a performance baseline with a simple architecture. |
+| [02-CNN.ipynb](notebooks/02-CNN.ipynb) | **Convolutional Neural Network** - Introduction to CNNs with convolutional and pooling layers using `nn.Sequential`. Demonstrates the advantage of CNNs over DNNs for image tasks. |
+| [03-RGB-CNN.ipynb](notebooks/03-RGB-CNN.ipynb) | **RGB CNN** - CNN classifier that utilizes full RGB color information instead of grayscale, improving feature extraction from color images. |
+| [04-optimized-CNN.ipynb](notebooks/04-optimized-CNN.ipynb) | **Hyperparameter Optimization** - Uses Optuna for automated hyperparameter tuning to find optimal network architecture and training parameters. |
+| [05-augmented-CNN.ipynb](notebooks/05-augmented-CNN.ipynb) | **Data Augmentation** - Trains the optimized CNN architecture with image augmentation techniques for improved generalization and robustness. |
+
+## Requirements
+
+- Python >=3.10, <3.13
+- PyTorch >=2.0
+- torchvision >=0.15
+- numpy >=1.24
+
+## License
+
+This project is licensed under the GPLv3 License - see the [LICENSE](LICENSE) file for details.

cifar10_tools-0.4.0/README.md

@@ -0,0 +1,38 @@
+# PyTorch: CIFAR-10 Demonstration
+
+A progressive deep learning tutorial for image classification on the CIFAR-10 dataset using PyTorch. This project demonstrates the evolution from basic deep neural networks to optimized convolutional neural networks with data augmentation. It also provides a set of utility functions as a PyPI package for use in other projects.
+
+[View on PyPI](https://pypi.org/project/cifar10_tools)
+
+## Installation
+
+Install the helper tools package locally in editable mode:
+
+```bash
+pip install -e .
+```
+
+## Project Overview
+
+This repository contains a series of Jupyter notebooks that progressively build more sophisticated neural network architectures for the CIFAR-10 image classification task. Each notebook builds upon concepts from the previous one, demonstrating key deep learning techniques.
+
+## Notebooks
+
+| Notebook | Description |
+|----------|-------------|
+| [01-DNN.ipynb](notebooks/01-DNN.ipynb) | **Deep Neural Network** - Baseline fully-connected DNN classifier using `nn.Sequential`. Establishes a performance baseline with a simple architecture. |
+| [02-CNN.ipynb](notebooks/02-CNN.ipynb) | **Convolutional Neural Network** - Introduction to CNNs with convolutional and pooling layers using `nn.Sequential`. Demonstrates the advantage of CNNs over DNNs for image tasks. |
+| [03-RGB-CNN.ipynb](notebooks/03-RGB-CNN.ipynb) | **RGB CNN** - CNN classifier that utilizes full RGB color information instead of grayscale, improving feature extraction from color images. |
+| [04-optimized-CNN.ipynb](notebooks/04-optimized-CNN.ipynb) | **Hyperparameter Optimization** - Uses Optuna for automated hyperparameter tuning to find optimal network architecture and training parameters. |
+| [05-augmented-CNN.ipynb](notebooks/05-augmented-CNN.ipynb) | **Data Augmentation** - Trains the optimized CNN architecture with image augmentation techniques for improved generalization and robustness. |
+
+## Requirements
+
+- Python >=3.10, <3.13
+- PyTorch >=2.0
+- torchvision >=0.15
+- numpy >=1.24
+
+## License
+
+This project is licensed under the GPLv3 License - see the [LICENSE](LICENSE) file for details.

{cifar10_tools-0.2.0 → cifar10_tools-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "cifar10_tools"
-version = "0.2.0"
+version = "0.4.0"
 description = "Tools for training neural networks on the CIFAR-10 task with PyTorch and TensorFlow"
 authors = ["gperdrizet <george@perdrizet.org>"]
 readme = "README.md"

cifar10_tools-0.4.0/src/cifar10_tools/pytorch/__init__.py

@@ -0,0 +1,33 @@
+'''PyTorch utilities for CIFAR-10 classification.'''
+
+from cifar10_tools.pytorch.data import download_cifar10_data
+from cifar10_tools.pytorch.evaluation import evaluate_model
+from cifar10_tools.pytorch.training import train_model
+from cifar10_tools.pytorch.plotting import (
+    plot_sample_images,
+    plot_learning_curves,
+    plot_confusion_matrix,
+    plot_class_probability_distributions,
+    plot_evaluation_curves,
+    plot_optimization_results
+)
+from cifar10_tools.pytorch.hyperparameter_optimization import (
+    create_cnn,
+    train_trial,
+    create_objective
+)
+
+__all__ = [
+    'download_cifar10_data',
+    'evaluate_model',
+    'train_model',
+    'plot_sample_images',
+    'plot_learning_curves',
+    'plot_confusion_matrix',
+    'plot_class_probability_distributions',
+    'plot_evaluation_curves',
+    'plot_optimization_results',
+    'create_cnn',
+    'train_trial',
+    'create_objective'
+]

cifar10_tools-0.4.0/src/cifar10_tools/pytorch/data.py

@@ -0,0 +1,135 @@
+'''Data loading and preprocessing functions for CIFAR-10 dataset.'''
+
+from pathlib import Path
+import torch
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+
+def download_cifar10_data(data_dir: str='data/pytorch/cifar10'):
+    '''Download CIFAR-10 dataset using torchvision.datasets.'''
+
+    data_dir = Path(data_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    _ = datasets.CIFAR10(
+        root=data_dir,
+        train=True,
+        download=True
+    )
+
+    _ = datasets.CIFAR10(
+        root=data_dir,
+        train=False,
+        download=True
+    )
+
+
+def make_data_loaders(
+    data_dir: Path,
+    batch_size: int,
+    train_transform: transforms.Compose,
+    eval_transform: transforms.Compose,
+    device: torch.device | None = None,
+    download: bool = False,
+):
+    """
+    Loads CIFAR-10, applies preprocessing with separate train/eval transforms,
+    and returns DataLoaders.
+    
+    Args:
+        data_dir: Path to CIFAR-10 data directory
+        batch_size: Batch size for DataLoaders
+        train_transform: Transform to apply to training data
+        eval_transform: Transform to apply to validation and test data
+        device: Device to preload tensors onto. If None, data stays on CPU 
+                and transforms are applied on-the-fly during iteration.
+        download: Whether to download the dataset if not present
+        
+    Returns:
+        Tuple of (train_loader, val_loader, test_loader)
+    """
+
+    # Load datasets with respective transforms
+    train_dataset_full = datasets.CIFAR10(
+        root=data_dir,
+        train=True,
+        download=download,
+        transform=train_transform,
+    )
+    
+    val_test_dataset_full = datasets.CIFAR10(
+        root=data_dir,
+        train=True,
+        download=download,
+        transform=eval_transform,
+    )
+    
+    test_dataset = datasets.CIFAR10(
+        root=data_dir,
+        train=False,
+        download=download,
+        transform=eval_transform,
+    )
+
+    if device is not None:
+        # Preload entire dataset to device for faster training
+        X_train_full = torch.stack([img for img, _ in train_dataset_full]).to(device)
+        y_train_full = torch.tensor([label for _, label in train_dataset_full]).to(device)
+        
+        X_val_test_full = torch.stack([img for img, _ in val_test_dataset_full]).to(device)
+        y_val_test_full = torch.tensor([label for _, label in val_test_dataset_full]).to(device)
+
+        X_test = torch.stack([img for img, _ in test_dataset]).to(device)
+        y_test = torch.tensor([label for _, label in test_dataset]).to(device)
+
+        # Train/val split (80/20)
+        n_train = int(0.8 * len(X_train_full))
+        indices = torch.randperm(len(X_train_full))
+
+        X_train = X_train_full[indices[:n_train]]
+        y_train = y_train_full[indices[:n_train]]
+        X_val = X_val_test_full[indices[n_train:]]
+        y_val = y_val_test_full[indices[n_train:]]
+
+        # TensorDatasets
+        train_tensor_dataset = torch.utils.data.TensorDataset(X_train, y_train)
+        val_tensor_dataset = torch.utils.data.TensorDataset(X_val, y_val)
+        test_tensor_dataset = torch.utils.data.TensorDataset(X_test, y_test)
+        
+    else:
+        # Don't preload - use datasets directly for on-the-fly transforms
+        # Train/val split (80/20) using Subset
+        n_train = int(0.8 * len(train_dataset_full))
+        indices = torch.randperm(len(train_dataset_full)).tolist()
+        
+        train_indices = indices[:n_train]
+        val_indices = indices[n_train:]
+        
+        train_tensor_dataset = torch.utils.data.Subset(train_dataset_full, train_indices)
+        val_tensor_dataset = torch.utils.data.Subset(val_test_dataset_full, val_indices)
+        test_tensor_dataset = test_dataset
+
+    # DataLoaders
+    train_loader = DataLoader(
+        train_tensor_dataset,
+        batch_size=batch_size,
+        shuffle=True,
+    )
+    val_loader = DataLoader(
+        val_tensor_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+    )
+    test_loader = DataLoader(
+        test_tensor_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+    )
+
+    return train_loader, val_loader, test_loader
+
+
+if __name__ == '__main__':
+
+    download_cifar10_data()

cifar10_tools-0.4.0/src/cifar10_tools/pytorch/hyperparameter_optimization.py

@@ -0,0 +1,278 @@
+'''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 cifar10_tools.pytorch.data import make_data_loaders
+
+
+def create_cnn(
+    n_conv_blocks: int,
+    initial_filters: int,
+    n_fc_layers: int,
+    base_kernel_size: int,
+    conv_dropout_rate: float,
+    fc_dropout_rate: float,
+    pooling_strategy: str,
+    use_batch_norm: bool,
+    num_classes: int = 10,
+    in_channels: int = 3,
+    input_size: int = 32
+) -> nn.Sequential:
+    '''Create a CNN with configurable architecture.
+    
+    Args:
+        n_conv_blocks: Number of convolutional blocks (1-5)
+        initial_filters: Number of filters in first conv layer (doubles each block)
+        n_fc_layers: Number of fully connected layers (1-8)
+        base_kernel_size: Base kernel size (decreases by 2 per block, min 3)
+        conv_dropout_rate: Dropout probability after convolutional blocks
+        fc_dropout_rate: Dropout probability in fully connected layers
+        pooling_strategy: Pooling type ('max' or 'avg')
+        use_batch_norm: Whether to use batch normalization
+        num_classes: Number of output classes (default: 10 for CIFAR-10)
+        in_channels: Number of input channels (default: 3 for RGB)
+        input_size: Input image size (default: 32 for CIFAR-10)
+    
+    Returns:
+        nn.Sequential model
+    '''
+    layers = []
+    current_channels = in_channels
+    current_size = input_size
+    
+    # Convolutional blocks
+    for block_idx in range(n_conv_blocks):
+        out_channels = initial_filters * (2 ** block_idx)
+        kernel_size = max(3, base_kernel_size - 2 * block_idx)
+        padding = kernel_size // 2
+        
+        # First conv in block
+        layers.append(nn.Conv2d(current_channels, out_channels, kernel_size=kernel_size, padding=padding))
+
+        if use_batch_norm:
+            layers.append(nn.BatchNorm2d(out_channels))
+
+        layers.append(nn.ReLU())
+        
+        # Second conv in block
+        layers.append(nn.Conv2d(out_channels, out_channels, kernel_size=kernel_size, padding=padding))
+
+        if use_batch_norm:
+            layers.append(nn.BatchNorm2d(out_channels))
+
+        layers.append(nn.ReLU())
+        
+        # Pooling
+        if pooling_strategy == 'max':
+            layers.append(nn.MaxPool2d(2, 2))
+        else:  # avg
+            layers.append(nn.AvgPool2d(2, 2))
+        
+        layers.append(nn.Dropout(conv_dropout_rate))
+        
+        current_channels = out_channels
+        current_size //= 2
+    
+    # Calculate flattened size
+    final_channels = initial_filters * (2 ** (n_conv_blocks - 1))
+    flattened_size = final_channels * current_size * current_size
+    
+    # Classifier - dynamic FC layers with halving pattern
+    layers.append(nn.Flatten())
+    
+    # Generate FC layer sizes by halving from flattened_size
+    fc_sizes = []
+    current_fc_size = flattened_size // 2
+    for _ in range(n_fc_layers):
+        fc_sizes.append(max(10, current_fc_size))  # Minimum 10 units
+        current_fc_size //= 2
+    
+    # Add FC layers
+    in_features = flattened_size
+    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,
+    train_transform,
+    eval_transform,
+    n_epochs: int,
+    device: torch.device,
+    num_classes: int = 10,
+    in_channels: int = 3
+) -> 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 CIFAR-10 data
+        train_transform: Transform to apply to training data
+        eval_transform: Transform to apply to validation data
+        n_epochs: Number of epochs per trial
+        device: Device to train on (cuda or cpu)
+        num_classes: Number of output classes (default: 10)
+        in_channels: Number of input channels (default: 3 for RGB)
+    
+    Returns:
+        Objective function for optuna.Study.optimize()
+    
+    Example:
+        >>> objective = create_objective(data_dir, transform, transform, n_epochs=50, device=device)
+        >>> study = optuna.create_study(direction='maximize')
+        >>> study.optimize(objective, n_trials=100)
+    '''
+    
+    def objective(trial: optuna.Trial) -> float:
+        '''Optuna objective function for CNN hyperparameter optimization.'''
+        
+        # Suggest hyperparameters
+        batch_size = trial.suggest_categorical('batch_size', [64, 128, 256, 512, 1024])
+        n_conv_blocks = trial.suggest_int('n_conv_blocks', 1, 5)
+        initial_filters = trial.suggest_categorical('initial_filters', [8, 16, 32, 64, 128])
+        n_fc_layers = trial.suggest_int('n_fc_layers', 1, 8)
+        base_kernel_size = trial.suggest_int('base_kernel_size', 3, 7)
+        conv_dropout_rate = trial.suggest_float('conv_dropout_rate', 0.0, 0.5)
+        fc_dropout_rate = trial.suggest_float('fc_dropout_rate', 0.2, 0.75)
+        pooling_strategy = trial.suggest_categorical('pooling_strategy', ['max', 'avg'])
+        use_batch_norm = trial.suggest_categorical('use_batch_norm', [True, False])
+        learning_rate = trial.suggest_float('learning_rate', 1e-5, 1e-1, log=True)
+        optimizer_name = trial.suggest_categorical('optimizer', ['Adam', 'SGD', 'RMSprop'])
+        
+        # Create data loaders with suggested batch size
+        train_loader, val_loader, _ = make_data_loaders(
+            data_dir=data_dir,
+            batch_size=batch_size,
+            train_transform=train_transform,
+            eval_transform=eval_transform,
+            device=device,
+            download=False
+        )
+        
+        # Create model
+        model = create_cnn(
+            n_conv_blocks=n_conv_blocks,
+            initial_filters=initial_filters,
+            n_fc_layers=n_fc_layers,
+            base_kernel_size=base_kernel_size,
+            conv_dropout_rate=conv_dropout_rate,
+            fc_dropout_rate=fc_dropout_rate,
+            pooling_strategy=pooling_strategy,
+            use_batch_norm=use_batch_norm,
+            num_classes=num_classes,
+            in_channels=in_channels
+        ).to(device)
+        
+        # Define optimizer
+        if optimizer_name == 'Adam':
+            optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+        elif optimizer_name == 'SGD':
+            momentum = trial.suggest_float('sgd_momentum', 0.8, 0.99)
+            optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
+        
+        else:  # RMSprop
+            optimizer = optim.RMSprop(model.parameters(), lr=learning_rate)
+        
+        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 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

{cifar10_tools-0.2.0 → cifar10_tools-0.4.0}/src/cifar10_tools/pytorch/plotting.py

@@ -10,18 +10,18 @@ def plot_sample_images(
     class_names: list[str],
     nrows: int = 2,
     ncols: int = 5,
-    figsize: tuple[float, float] | None = None,
-    cmap: str = 'gray'
+    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).
-        cmap: Colormap for displaying images.
         
     Returns:
         Tuple of (figure, axes array).
@@ -36,11 +36,24 @@ def plot_sample_images(
         # Get image and label from dataset
         img, label = dataset[i]
         
-        # Unnormalize and squeeze for plotting
+        # Unnormalize for plotting
         img = img * 0.5 + 0.5
-        img = img.numpy().squeeze()
+        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.imshow(img, cmap=cmap)
         ax.axis('off')
 
     plt.tight_layout()
@@ -209,7 +222,7 @@ def plot_evaluation_curves(
         roc_auc = auc(fpr, tpr)
         ax1.plot(fpr, tpr, label=class_name)
 
-    ax1.plot([0, 1], [0, 1], 'k--', label='Random classifier')
+    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)
@@ -235,4 +248,60 @@ def plot_evaluation_curves(
 
     plt.tight_layout()
     
-    return fig, (ax1, ax2)
+    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

{cifar10_tools-0.2.0 → cifar10_tools-0.4.0}/src/cifar10_tools/pytorch/training.py

@@ -12,11 +12,28 @@ def train_model(
     criterion: nn.Module,
     optimizer: optim.Optimizer,
     epochs: int = 10,
-    print_every: int = 1
+    print_every: int = 1,
+    device: torch.device | str | None = None
 ) -> dict[str, list[float]]:
     '''Training loop for PyTorch classification model.
     
-    Note: Assumes data is already on the correct device.
+    Handles both pre-loaded GPU data and lazy-loading (CPU data moved per-batch).
+    
+    Args:
+        model: PyTorch model to train.
+        train_loader: DataLoader for training data.
+        val_loader: DataLoader for validation data.
+        criterion: Loss function.
+        optimizer: Optimizer.
+        epochs: Number of training epochs.
+        print_every: Print progress every n epochs.
+        device: Device to move data to per-batch. If None, assumes data is 
+                already on the correct device (GPU pre-loading). If specified,
+                data will be moved to this device per-batch (lazy loading).
+    
+    Returns:
+        Dictionary containing training history with keys:
+        'train_loss', 'val_loss', 'train_accuracy', 'val_accuracy'.
     '''
 
     history = {'train_loss': [], 'val_loss': [], 'train_accuracy': [], 'val_accuracy': []}
@@ -30,6 +47,11 @@ def train_model(
         total = 0
 
         for images, labels in train_loader:
+            
+            # Move to device if lazy loading
+            if device is not None:
+                images = images.to(device, non_blocking=True)
+                labels = labels.to(device, non_blocking=True)
 
             # Forward pass
             optimizer.zero_grad()
@@ -59,6 +81,11 @@ def train_model(
         with torch.no_grad():
 
             for images, labels in val_loader:
+                
+                # Move to device if lazy loading
+                if device is not None:
+                    images = images.to(device, non_blocking=True)
+                    labels = labels.to(device, non_blocking=True)
 
                 outputs = model(images)
                 loss = criterion(outputs, labels)
@@ -88,6 +115,4 @@ def train_model(
                 f'val_accuracy: {val_accuracy:.2f}%'
             )
 
-    print('\nTraining complete.')
-
     return history

cifar10_tools-0.2.0/PKG-INFO

@@ -1,35 +0,0 @@
-Metadata-Version: 2.4
-Name: cifar10_tools
-Version: 0.2.0
-Summary: Tools for training neural networks on the CIFAR-10 task with PyTorch and TensorFlow
-License: GPLv3
-License-File: LICENSE
-Keywords: Python,Machine learning,Deep learning,CNNs,Computer vision,Image classification,CIFAR-10
-Author: gperdrizet
-Author-email: george@perdrizet.org
-Requires-Python: >=3.10,<3.13
-Classifier: Development Status :: 3 - Alpha
-Classifier: Intended Audience :: Developers
-Classifier: Intended Audience :: Education
-Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
-Classifier: License :: Other/Proprietary License
-Classifier: Operating System :: OS Independent
-Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
-Classifier: Topic :: Scientific/Engineering :: Image Recognition
-Provides-Extra: tensorflow
-Requires-Dist: numpy (>=1.24)
-Requires-Dist: torch (>=2.0)
-Requires-Dist: torchvision (>=0.15)
-Project-URL: Documentation, https://gperdrizet.github.io/CIFAR10/README.md
-Project-URL: Homepage, https://github.com/gperdrizet/CIFAR10
-Project-URL: Issues, https://github.com/gperdrizet/CIFAR10/issues
-Project-URL: PyPI, https://pypi.org/project/cifar10_tools
-Project-URL: Repository, https://github.com/gperdrizet/CIFAR10
-Description-Content-Type: text/markdown
-
-# PyTorch: CIFAR10 demonstration

cifar10_tools-0.2.0/README.md

@@ -1 +0,0 @@
-# PyTorch: CIFAR10 demonstration

cifar10_tools-0.2.0/src/cifar10_tools/pytorch/data.py

@@ -1,27 +0,0 @@
-'''Data download function for CIFAR-10 dataset. Use to pre-download data
-during devcontainer creation'''
-
-from pathlib import Path
-from torchvision import datasets
-
-def download_cifar10_data(data_dir: str='data/pytorch/cifar10'):
-    '''Download CIFAR-10 dataset using torchvision.datasets.'''
-
-    data_dir = Path(data_dir)
-    data_dir.mkdir(parents=True, exist_ok=True)
-
-    _ = datasets.CIFAR10(
-        root=data_dir,
-        train=True,
-        download=True
-    )
-
-    _ = datasets.CIFAR10(
-        root=data_dir,
-        train=False,
-        download=True
-    )
-
-if __name__ == '__main__':
-
-    download_cifar10_data()