pytorch-kito 0.2.0__py3-none-any.whl

kito/data/__init__.py

@@ -0,0 +1,35 @@
+# src/kito/data/__init__.py
+"""Kito Data - Datasets, preprocessing, and pipelines"""
+
+from kito.data.datasets import KitoDataset, H5Dataset, MemDataset
+from kito.data.preprocessing import (
+    Preprocessing,
+    Pipeline,
+    Normalize,
+    Standardization,
+    ToTensor
+)
+from kito.data.datapipeline import GenericDataPipeline, BaseDataPipeline
+from kito.data.registry import DATASETS, PREPROCESSING
+
+__all__ = [
+    # Datasets
+    "KitoDataset",
+    "H5Dataset",
+    "MemDataset",
+
+    # Preprocessing
+    "Preprocessing",
+    "Pipeline",
+    "Normalize",
+    "Standardization",
+    "ToTensor",
+
+    # Pipelines
+    "GenericDataPipeline",
+    "BaseDataPipeline",
+
+    # Registries
+    "DATASETS",
+    "PREPROCESSING",
+]

kito/data/datapipeline.py

@@ -0,0 +1,273 @@
+"""
+DataPipeline classes - Orchestrate data pipeline.
+
+DataPipeline encapsulates:
+- Dataset creation
+- Preprocessing setup
+- Train/val/test splitting
+- DataLoader configuration
+
+This separates data concerns from model training logic.
+"""
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from torch.utils.data import DataLoader, Subset, DistributedSampler
+
+from kito.config.moduleconfig import KitoModuleConfig
+from kito.data.datasets import KitoDataset
+from kito.data.preprocessed_dataset import PreprocessedDataset
+from kito.data.preprocessing import Preprocessing
+
+
+class BaseDataPipeline(ABC):
+    """
+    Base class for data modules.
+
+    A DataModule encapsulates all data-related logic:
+    - How to load data (dataset)
+    - How to preprocess data
+    - How to split into train/val/test
+    - How to create DataLoaders
+
+    Usage:
+        class MyDataModule(BaseDataModule):
+            def setup(self):
+                self.dataset = MyDataset(self.data_config.dataset_path)
+                # ... setup preprocessing, splits, etc.
+
+            def train_dataloader(self):
+                return DataLoader(self.train_dataset, ...)
+
+        # In training
+        dm = MyDataModule(config)
+        dm.setup()
+        engine.fit(datamodule=dm)
+    """
+
+    def __init__(self, config: KitoModuleConfig):
+        """
+        Initialize DataModule.
+
+        Args:
+            config: Full module configuration (contains data config)
+        """
+        self.config = config
+        self.data_config = config.data
+
+        # Will be set by setup()
+        self.dataset = None
+        self.train_dataset = None
+        self.val_dataset = None
+        self.test_dataset = None
+
+        self.train_loader = None
+        self.val_loader = None
+        self.test_loader = None
+
+    @abstractmethod
+    def setup(self):
+        """
+        Setup datasets and preprocessing.
+
+        This method should:
+        1. Load raw dataset
+        2. Apply preprocessing (if any)
+        3. Create train/val/test splits
+        4. Create DataLoaders
+
+        Called once before training.
+        """
+        pass
+
+    @abstractmethod
+    def train_dataloader(self) -> DataLoader:
+        """Return training DataLoader."""
+        pass
+
+    @abstractmethod
+    def val_dataloader(self) -> DataLoader:
+        """Return validation DataLoader."""
+        pass
+
+    @abstractmethod
+    def test_dataloader(self) -> DataLoader:
+        """Return test DataLoader."""
+        pass
+
+
+class GenericDataPipeline(BaseDataPipeline):
+    """
+    Generic DataModule that works with any registered dataset.
+
+    This is the standard DataModule used by DataModuleFactory.
+    It handles:
+    - Loading dataset from registry
+    - Applying preprocessing pipeline
+    - Creating train/val/test splits
+    - Setting up DataLoaders with DDP support
+
+    Args:
+        config: Full module configuration
+        dataset: Pre-instantiated dataset (optional)
+        preprocessing: Pre-instantiated preprocessing (optional)
+
+    Example:
+        >>> dm = GenericDataModule(config)
+        >>> dm.setup()
+        >>> train_loader = dm.train_dataloader()
+    """
+
+    def __init__(
+        self,
+        config: KitoModuleConfig,
+        dataset: Optional[KitoDataset] = None,
+        preprocessing: Optional[Preprocessing] = None
+    ):
+        super().__init__(config)
+        self._dataset = dataset  # Pre-instantiated dataset
+        self._preprocessing = preprocessing  # Pre-instantiated preprocessing
+
+    def setup(self):
+        """
+        Setup data pipeline.
+
+        Steps:
+        1. Use pre-instantiated dataset or load from registry
+        2. Wrap with preprocessing if provided
+        3. Create train/val/test splits
+        4. Create DataLoaders
+        """
+        # 1. Get dataset
+        if self._dataset is not None:
+            raw_dataset = self._dataset
+        else:
+            raise ValueError("Dataset not provided to GenericDataModule")
+
+        # 2. Apply preprocessing
+        if self._preprocessing is not None:
+            self.dataset = PreprocessedDataset(raw_dataset, self._preprocessing)
+        else:
+            self.dataset = raw_dataset
+
+        # 3. Create splits
+        self._create_splits()
+
+        # 4. Create DataLoaders
+        self._create_dataloaders()
+
+    def _create_splits(self):
+        """
+        Create train/val/test splits using Subset.
+
+        Splits based on data_config.train_ratio and data_config.val_ratio.
+        """
+        total_samples = self.data_config.total_samples
+        if total_samples is None:
+            total_samples = len(self.dataset)
+
+        # Calculate split indices
+        train_size = int(total_samples * self.data_config.train_ratio)
+        val_size = int(total_samples * self.data_config.val_ratio)
+        test_size = total_samples - train_size - val_size
+
+        # Create subsets
+        self.train_dataset = Subset(
+            self.dataset,
+            list(range(0, train_size))
+        )
+        self.val_dataset = Subset(
+            self.dataset,
+            list(range(train_size, train_size + val_size))
+        )
+        self.test_dataset = Subset(
+            self.dataset,
+            list(range(train_size + val_size, total_samples))
+        )
+
+    def _create_dataloaders(self):
+        """
+        Create DataLoaders with proper settings.
+
+        Handles:
+        - Distributed training (DistributedSampler)
+        - num_workers, pin_memory, etc.
+        - Batch size from config
+        """
+        batch_size = self.config.training.batch_size
+        distributed = self.config.training.distributed_training
+
+        # DataLoader kwargs
+        loader_kwargs = {
+            'batch_size': batch_size,
+            'num_workers': self.data_config.num_workers,
+            'pin_memory': self.data_config.pin_memory,
+            'persistent_workers': self.data_config.persistent_workers if self.data_config.num_workers > 0 else False,
+            'prefetch_factor': self.data_config.prefetch_factor if self.data_config.num_workers > 0 else None,
+        }
+
+        # Train loader (with shuffling)
+        if distributed:
+            train_sampler = DistributedSampler(self.train_dataset, shuffle=True)
+            self.train_loader = DataLoader(
+                self.train_dataset,
+                sampler=train_sampler,
+                shuffle=False,  # Don't shuffle when using sampler
+                **loader_kwargs
+            )
+        else:
+            self.train_loader = DataLoader(
+                self.train_dataset,
+                shuffle=True,
+                **loader_kwargs
+            )
+
+        # Val loader (no shuffling)
+        if distributed:
+            val_sampler = DistributedSampler(self.val_dataset, shuffle=False)
+            self.val_loader = DataLoader(
+                self.val_dataset,
+                sampler=val_sampler,
+                shuffle=False,
+                **loader_kwargs
+            )
+        else:
+            self.val_loader = DataLoader(
+                self.val_dataset,
+                shuffle=False,
+                **loader_kwargs
+            )
+
+        # Test loader (no shuffling)
+        if distributed:
+            test_sampler = DistributedSampler(self.test_dataset, shuffle=False)
+            self.test_loader = DataLoader(
+                self.test_dataset,
+                sampler=test_sampler,
+                shuffle=False,
+                **loader_kwargs
+            )
+        else:
+            self.test_loader = DataLoader(
+                self.test_dataset,
+                shuffle=False,
+                **loader_kwargs
+            )
+
+    def train_dataloader(self) -> DataLoader:
+        """Return training DataLoader."""
+        if self.train_loader is None:
+            raise RuntimeError("Call setup() before accessing dataloaders")
+        return self.train_loader
+
+    def val_dataloader(self) -> DataLoader:
+        """Return validation DataLoader."""
+        if self.val_loader is None:
+            raise RuntimeError("Call setup() before accessing dataloaders")
+        return self.val_loader
+
+    def test_dataloader(self) -> DataLoader:
+        """Return test DataLoader."""
+        if self.test_loader is None:
+            raise RuntimeError("Call setup() before accessing dataloaders")
+        return self.test_loader

kito/data/datasets.py

@@ -0,0 +1,166 @@
+import h5py
+
+from torch.utils.data import Dataset
+from abc import ABC, abstractmethod
+
+from kito.data.registry import DATASETS
+
+
+class KitoDataset(Dataset, ABC):
+    """
+    Abstract base class for all datasets.
+
+    Provides common data loading and preprocessing pattern.
+    Subclasses must implement _load_sample() to define how data is loaded.
+
+    The standard workflow is:
+    1. _load_sample(index) - Load raw data (subclass implements)
+    2. Return data and labels
+
+    Note: Preprocessing is now handled by PreprocessedDataset wrapper,
+    so _preprocess_data() is removed from here.
+    """
+
+    @abstractmethod
+    def _load_sample(self, index):
+        """
+        Load a single raw sample from the data source.
+
+        Args:
+            index: Sample index
+
+        Returns:
+            Tuple of (data, labels) as numpy arrays or tensors
+        """
+        pass
+
+    def __getitem__(self, index):
+        """
+        Get a single sample from the dataset.
+
+        Standard workflow:
+        1. Load raw sample
+        2. Return (preprocessing happens in PreprocessedDataset wrapper)
+        """
+        data, labels = self._load_sample(index)
+        return data, labels
+
+    @abstractmethod
+    def __len__(self):
+        """Get the total number of samples in the dataset."""
+        pass
+
+
+@DATASETS.register('h5dataset')
+class H5Dataset(KitoDataset):
+    """
+    HDF5 dataset for PyTorch with lazy loading.
+
+    Loads 'data' and 'labels' from HDF5 file with lazy loading for
+    multiprocessing compatibility (DataLoader with num_workers > 0).
+
+    Args:
+        path: Path to HDF5 file
+
+    HDF5 Structure Expected:
+        - 'data': Input data array (N, ...)
+        - 'labels': Target labels array (N, ...)
+
+    Example:
+        >>> dataset = H5Dataset("train.h5")
+        >>> data, labels = dataset[0]
+        >>> 
+        >>> # Works with DataLoader
+        >>> loader = DataLoader(dataset, batch_size=32, num_workers=4)
+    """
+
+    def __init__(self, path: str):
+        self.file_path = path
+
+        # Lazy-loaded attributes (set in _lazy_load)
+        self.dataset_data = None
+        self.dataset_labels = None
+        self.h5file = None
+
+        # Get dataset length (without lazy loading)
+        with h5py.File(self.file_path, 'r') as file:
+            self.dataset_len = len(file["data"])
+
+    def _lazy_load(self):
+        """
+        Open HDF5 file and get dataset references.
+        Called automatically in _load_sample().
+        """
+        if self.dataset_data is None or self.dataset_labels is None:
+            try:
+                self.h5file = h5py.File(self.file_path, 'r')
+                self.dataset_data = self.h5file["data"]
+                self.dataset_labels = self.h5file["labels"]
+            except (OSError, KeyError) as e:
+                raise RuntimeError(f"Failed to load H5 file '{self.file_path}': {e}")
+
+    def _load_sample(self, index):
+        """Load sample from HDF5 file with lazy loading."""
+        self._lazy_load()
+        return self.dataset_data[index], self.dataset_labels[index]
+
+    def __len__(self):
+        return self.dataset_len
+
+    def __del__(self):
+        """Close HDF5 file when object is destroyed."""
+        if hasattr(self, 'h5file') and self.h5file is not None:
+            self.h5file.close()
+
+    def __getstate__(self):
+        """
+        Prepare object for pickling (needed for DataLoader with num_workers > 0).
+        Remove non-picklable HDF5 file handles.
+        """
+        state = self.__dict__.copy()
+        # Remove HDF5 file handle and dataset references
+        state['dataset_data'] = None
+        state['dataset_labels'] = None
+        if 'h5file' in state:
+            del state['h5file']
+        return state
+
+    def __setstate__(self, state):
+        """Restore state after unpickling. File handles will be reloaded lazily."""
+        self.__dict__.update(state)
+
+
+@DATASETS.register('memdataset')
+class MemDataset(KitoDataset):
+    """
+    In-memory dataset for PyTorch.
+
+    Stores data and labels in memory (as numpy arrays or tensors).
+    Useful when dataset fits in RAM for faster training.
+
+    Args:
+        x: Input data array (N, ...)
+        y: Target labels array (N, ...)
+
+    Example:
+        >>> import numpy as np
+        >>> x = np.random.randn(100, 10, 64, 64, 1)
+        >>> y = np.random.randn(100, 10, 64, 64, 1)
+        >>> dataset = MemDataset(x, y)
+        >>> data, labels = dataset[0]
+    """
+
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
+
+        # Validate shapes
+        if len(x) != len(y):
+            raise ValueError(f"x and y must have same length. Got {len(x)} and {len(y)}")
+
+    def _load_sample(self, index):
+        """Load sample from memory arrays."""
+        return self.x[index], self.y[index]
+
+    def __len__(self):
+        return self.x.shape[0]

kito/data/preprocessed_dataset.py

@@ -0,0 +1,57 @@
+"""
+PreprocessedDataset - Wraps a dataset and applies preprocessing.
+
+Separates data loading from preprocessing for flexibility.
+
+Example:
+    raw_dataset = H5Dataset('data.h5')
+    preprocessing = Pipeline([Detrend(), Standardization()])
+    dataset = PreprocessedDataset(raw_dataset, preprocessing)
+"""
+from torch.utils.data import Dataset
+
+
+class PreprocessedDataset(Dataset):
+    """
+    Wraps a base dataset and applies preprocessing on-the-fly.
+
+    This allows:
+    - Keeping datasets "dumb" (just load data)
+    - Composable preprocessing
+    - Easy experimentation (swap preprocessing without changing dataset)
+
+    Args:
+        base_dataset: Underlying dataset (H5Dataset, MemDataset, etc.)
+        preprocessing: Preprocessing instance or None
+
+    Example:
+        >>> raw_dataset = H5Dataset('data.h5')
+        >>> preprocessing = Standardization(mean=0, std=1)
+        >>> dataset = PreprocessedDataset(raw_dataset, preprocessing)
+        >>> data, labels = dataset[0]  # Automatically preprocessed
+    """
+
+    def __init__(self, base_dataset: Dataset, preprocessing=None):
+        self.base_dataset = base_dataset
+        self.preprocessing = preprocessing
+
+    def __getitem__(self, index):
+        # Load raw data
+        data, labels = self.base_dataset[index]
+
+        # Apply preprocessing if specified
+        if self.preprocessing is not None:
+            data, labels = self.preprocessing(data, labels)
+
+        return data, labels
+
+    def __len__(self):
+        return len(self.base_dataset)
+
+    def __repr__(self):
+        return (
+            f"PreprocessedDataset(\n"
+            f"  base_dataset={self.base_dataset},\n"
+            f"  preprocessing={self.preprocessing}\n"
+            f")"
+        )