pytorch-kito 0.2.0__py3-none-any.whl

kito/__init__.py

@@ -0,0 +1,49 @@
+"""
+Kito: Effortless PyTorch Training
+
+Define your model, Kito handles the rest.
+"""
+
+__version__ = "0.2.0"
+
+# Import main classes for top-level access
+from kito.engine import Engine
+from kito.module import KitoModule
+
+# Import common data classes
+from kito.data.datasets import H5Dataset, MemDataset, KitoDataset
+from kito.data.datapipeline import GenericDataPipeline
+from kito.data.preprocessing import (
+    Preprocessing,
+    Pipeline,
+    Normalize,
+    Standardization,
+    ToTensor
+)
+
+# Import registries
+from kito.data.registry import DATASETS, PREPROCESSING
+
+# Define what's available with "from kito import *"
+__all__ = [
+    # Core
+    "Engine",
+    "KitoModule",
+
+    # Data
+    "H5Dataset",
+    "MemDataset",
+    "KitoDataset",
+    "GenericDataPipeline",
+
+    # Preprocessing
+    "Preprocessing",
+    "Pipeline",
+    "Normalize",
+    "Standardization",
+    "ToTensor",
+
+    # Registries
+    "DATASETS",
+    "PREPROCESSING",
+]

kito/callbacks/__init__.py

@@ -0,0 +1,20 @@
+# src/kito/callbacks/__init__.py
+"""Kito Callbacks - For custom training behavior"""
+
+from kito.callbacks.callback_base import Callback, CallbackList
+from kito.callbacks.modelcheckpoint import ModelCheckpoint
+from kito.callbacks.csv_logger import CSVLogger
+from kito.callbacks.txt_logger import TextLogger
+from kito.callbacks.tensorboard_callbacks import TensorBoardScalars, TensorBoardGraph, TensorBoardHistograms
+# ... other callbacks
+
+__all__ = [
+    "Callback",
+    "CallbackList",
+    "ModelCheckpoint",
+    "CSVLogger",
+    "TextLogger",
+    "TensorBoardScalars",
+    "TensorBoardGraph",
+    "TensorBoardHistograms"
+]

kito/callbacks/callback_base.py

@@ -0,0 +1,107 @@
+"""
+Modern callback system for BaseModule.
+
+Inspired by Keras and PyTorch Lightning callback patterns.
+Each callback is independent and handles a single concern.
+"""
+from abc import ABC
+
+
+class Callback(ABC):
+    """
+    Base class for all callbacks.
+
+    Callbacks allow you to customize the training loop behavior
+    by hooking into specific events (epoch start/end, batch start/end, etc.).
+
+    All methods have default no-op implementations, so you only need to
+    override the ones you care about.
+    """
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Called at the beginning of training."""
+        pass
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Called at the end of training."""
+        pass
+
+    def on_epoch_begin(self, epoch, engine, model, **kwargs):
+        """Called at the beginning of each epoch."""
+        pass
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """
+        Called at the end of each epoch.
+
+        Args:
+            epoch: Current epoch number (1-indexed)
+            engine: Reference to the Engine
+            model: The PyTorch model
+            logs: Dictionary of metrics (e.g., {'train_loss': 0.5, 'val_loss': 0.3})
+            **kwargs: Additional context (val_data, val_outputs, etc.)
+        """
+        pass
+
+    def on_train_batch_begin(self, batch, engine, model, **kwargs):
+        """Called at the beginning of each training batch."""
+        pass
+
+    def on_train_batch_end(self, batch, engine, model, logs=None, **kwargs):
+        """Called at the end of each training batch."""
+        pass
+
+    def on_validation_begin(self, epoch, engine, model, **kwargs):
+        """Called at the beginning of validation."""
+        pass
+
+    def on_validation_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Called at the end of validation."""
+        pass
+
+
+class CallbackList:
+    """
+    Container for managing multiple callbacks.
+
+    Iterates through all callbacks and calls the appropriate method.
+    """
+
+    def __init__(self, callbacks=None):
+        self.callbacks = callbacks or []
+
+    def append(self, callback):
+        """Add a callback to the list."""
+        self.callbacks.append(callback)
+
+    def on_train_begin(self, engine, model, **kwargs):
+        for callback in self.callbacks:
+            callback.on_train_begin(engine, model, **kwargs)
+
+    def on_train_end(self, engine, model, **kwargs):
+        for callback in self.callbacks:
+            callback.on_train_end(engine, model, **kwargs)
+
+    def on_epoch_begin(self, epoch, engine, model, **kwargs):
+        for callback in self.callbacks:
+            callback.on_epoch_begin(epoch, engine, model, **kwargs)
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        for callback in self.callbacks:
+            callback.on_epoch_end(epoch, engine, model, logs, **kwargs)
+
+    def on_train_batch_begin(self, batch, engine, model, **kwargs):
+        for callback in self.callbacks:
+            callback.on_train_batch_begin(batch, engine, model, **kwargs)
+
+    def on_train_batch_end(self, batch, engine, model, logs=None, **kwargs):
+        for callback in self.callbacks:
+            callback.on_train_batch_end(batch, engine, model, logs, **kwargs)
+
+    def on_validation_begin(self, epoch, engine, model, **kwargs):
+        for callback in self.callbacks:
+            callback.on_validation_begin(epoch, engine, model, **kwargs)
+
+    def on_validation_end(self, epoch, engine, model, logs=None, **kwargs):
+        for callback in self.callbacks:
+            callback.on_validation_end(epoch, engine, model, logs, **kwargs)

kito/callbacks/csv_logger.py

@@ -0,0 +1,66 @@
+import csv
+import os
+
+from kito.callbacks.callback_base import Callback
+
+
+class CSVLogger(Callback):
+    """
+    Log training metrics to a CSV file.
+
+    Args:
+        filename: Path to CSV file
+        separator: Column separator (default: ',')
+        append: Append to existing file or overwrite
+
+    Example:
+        csv_logger = CSVLogger('logs/training_log.csv')
+    """
+
+    def __init__(
+            self,
+            filename: str,
+            separator: str = ',',
+            append: bool = False
+    ):
+        self.filename = filename
+        self.separator = separator
+        self.append = append
+        self.writer = None
+        self.file = None
+        self.keys = None
+
+        # Create directory
+        os.makedirs(os.path.dirname(filename) if os.path.dirname(filename) else '.', exist_ok=True)
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Open CSV file and write header."""
+        mode = 'a' if self.append else 'w'
+        self.file = open(self.filename, mode, newline='')
+        self.writer = csv.writer(self.file, delimiter=self.separator)
+
+        # If not appending, we'll write header on first epoch
+        if not self.append:
+            self.keys = None
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Write metrics to CSV."""
+        if logs is None:
+            return
+
+        # Add epoch to logs
+        row_dict = {'epoch': epoch, **logs}
+
+        # Write header if first time
+        if self.keys is None:
+            self.keys = list(row_dict.keys())
+            self.writer.writerow(self.keys)
+
+        # Write values
+        self.writer.writerow([row_dict.get(k, '') for k in self.keys])
+        self.file.flush()
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close CSV file."""
+        if self.file:
+            self.file.close()

kito/callbacks/ddp_aware_callback.py

@@ -0,0 +1,60 @@
+import torch.distributed as dist
+
+from kito.callbacks.callback_base import Callback
+
+
+class DDPAwareCallback(Callback):
+    """
+    Wrapper that makes any callback DDP-safe.
+
+    Only executes on rank 0 in distributed training.
+
+    Args:
+        callback: The callback to wrap
+
+    Example:
+        checkpoint = ModelCheckpoint('weights/best.pt')
+        ddp_checkpoint = DDPAwareCallback(checkpoint)
+    """
+
+    def __init__(self, callback: Callback):
+        self.callback = callback
+        self.is_driver = self._check_if_driver()
+
+    def _check_if_driver(self):
+        """Check if this is the driver process (rank 0)."""
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_rank() == 0
+        return True  # Single GPU or CPU
+
+    def on_train_begin(self, engine, model, **kwargs):
+        if self.is_driver:
+            self.callback.on_train_begin(engine, model, **kwargs)
+
+    def on_train_end(self, engine, model, **kwargs):
+        if self.is_driver:
+            self.callback.on_train_end(engine, model, **kwargs)
+
+    def on_epoch_begin(self, epoch, engine, model, **kwargs):
+        if self.is_driver:
+            self.callback.on_epoch_begin(epoch, engine, model, **kwargs)
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        if self.is_driver:
+            self.callback.on_epoch_end(epoch, engine, model, logs, **kwargs)
+
+    def on_train_batch_begin(self, batch, engine, model, **kwargs):
+        if self.is_driver:
+            self.callback.on_train_batch_begin(batch, engine, model, **kwargs)
+
+    def on_train_batch_end(self, batch, engine, model, logs=None, **kwargs):
+        if self.is_driver:
+            self.callback.on_train_batch_end(batch, engine, model, logs, **kwargs)
+
+    def on_validation_begin(self, epoch, engine, model, **kwargs):
+        if self.is_driver:
+            self.callback.on_validation_begin(epoch, engine, model, **kwargs)
+
+    def on_validation_end(self, epoch, engine, model, logs=None, **kwargs):
+        if self.is_driver:
+            self.callback.on_validation_end(epoch, engine, model, logs, **kwargs)

kito/callbacks/early_stopping_callback.py

@@ -0,0 +1,45 @@
+from kito.callbacks.callback_base import Callback
+
+
+class EarlyStoppingCallback(Callback):
+    """
+    Stop training when monitored metric stops improving.
+
+    Args:
+        monitor: Metric to monitor
+        patience: Number of epochs with no improvement before stopping
+        mode: 'min' or 'max'
+    """
+
+    def __init__(self, monitor='val_loss', patience=10, mode='min'):
+        self.monitor = monitor
+        self.patience = patience
+        self.mode = mode
+        self.best = float('inf') if mode == 'min' else float('-inf')
+        self.wait = 0
+        self.stopped_epoch = 0
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Check if training should stop."""
+        if logs is None:
+            return
+
+        current = logs.get(self.monitor)
+        if current is None:
+            return
+
+        # Check if improved
+        improved = (
+                (self.mode == 'min' and current < self.best) or
+                (self.mode == 'max' and current > self.best)
+        )
+
+        if improved:
+            self.best = current
+            self.wait = 0
+        else:
+            self.wait += 1
+            if self.wait >= self.patience:
+                self.stopped_epoch = epoch
+                engine.stop_training = True  # Requires Engine to support this
+                print(f"\nEarly stopping at epoch {epoch}")

kito/callbacks/modelcheckpoint.py

@@ -0,0 +1,78 @@
+import os
+import torch
+
+from kito.callbacks.callback_base import Callback
+
+
+class ModelCheckpoint(Callback):
+    """
+    Save model weights during training.
+
+    Args:
+        filepath: Path template for saving weights (can include {epoch}, {val_loss}, etc.)
+        monitor: Metric to monitor (e.g., 'val_loss')
+        save_best_only: Only save when monitored metric improves
+        mode: 'min' or 'max' depending on whether lower/higher is better
+        verbose: Print message when saving
+
+    Example:
+        checkpoint = ModelCheckpoint(
+            filepath='weights/model_epoch{epoch:02d}_valloss{val_loss:.4f}.pt',
+            monitor='val_loss',
+            save_best_only=True,
+            mode='min'
+        )
+    """
+
+    def __init__(
+            self,
+            filepath: str,
+            monitor: str = 'val_loss',
+            save_best_only: bool = True,
+            mode: str = 'min',
+            verbose: bool = False
+    ):
+        self.filepath = filepath
+        self.monitor = monitor
+        self.save_best_only = save_best_only
+        self.mode = mode
+        self.verbose = verbose
+
+        # Track best metric
+        self.best = float('inf') if mode == 'min' else float('-inf')
+
+        # Create directory if needed
+        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else '.', exist_ok=True)
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Save model if metric improved."""
+        if logs is None:
+            return
+
+        current = logs.get(self.monitor)
+        if current is None:
+            return
+
+        # Check if improved
+        improved = (
+                (self.mode == 'min' and current < self.best) or
+                (self.mode == 'max' and current > self.best)
+        )
+
+        if improved or not self.save_best_only:
+            self.best = current
+
+            # Format filepath
+            filepath = self.filepath.format(epoch=epoch, **logs)
+
+            # Save model (handle DDP)
+            if hasattr(model, 'module'):
+                state_dict = model.module.state_dict()
+            else:
+                state_dict = model.state_dict()
+
+            torch.save(state_dict, filepath)
+
+            if self.verbose:
+                print(f"\nEpoch {epoch}: {self.monitor} improved to {current:.4f}, "
+                      f"saving model to {filepath}")