pytorch-kito 0.2.0__py3-none-any.whl

kito/callbacks/tensorboard_callback_images.py

@@ -0,0 +1,298 @@
+"""
+Base and example callbacks for image plotting to TensorBoard.
+
+The BaseImagePlotter provides a foundation for custom image plotting callbacks.
+"""
+from abc import abstractmethod
+from typing import Optional, List
+
+import matplotlib.pyplot as plt
+from torch.utils.tensorboard import SummaryWriter
+
+from kito.callbacks.callback_base import Callback
+
+
+class BaseImagePlotter(Callback):
+    """
+    Base class for plotting images to TensorBoard.
+
+    Subclasses must implement `create_figure()` to define custom plotting logic.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+        tag: Tag for the image in TensorBoard
+        freq: Frequency of plotting (every N epochs)
+        batch_indices: Which batch indices to visualize (e.g., [0, 1, 2])
+
+    Example:
+        class MyCustomPlotter(BaseImagePlotter):
+            def create_figure(self, val_data, val_outputs, epoch, **kwargs):
+                fig, ax = plt.subplots()
+                # Your custom plotting logic here
+                ax.imshow(val_outputs[0].cpu().numpy())
+                return fig
+
+        plotter = MyCustomPlotter('logs/images', tag='predictions')
+    """
+
+    def __init__(
+            self,
+            log_dir: str,
+            tag: str = 'validation_images',
+            freq: int = 1,
+            batch_indices: Optional[List[int]] = None
+    ):
+        self.log_dir = log_dir
+        self.tag = tag
+        self.freq = freq
+        self.batch_indices = batch_indices or [0]
+        self.writer = None
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Initialize TensorBoard writer."""
+        self.writer = SummaryWriter(self.log_dir)
+
+    @abstractmethod
+    def create_figure(self, val_data, val_outputs, epoch, **kwargs):
+        """
+        Create matplotlib figure for visualization.
+
+        Must be implemented by subclasses.
+
+        Args:
+            val_data: Validation data (tuple of inputs and targets)
+            val_outputs: Model predictions on validation data
+            epoch: Current epoch number
+            **kwargs: Additional context from engine
+
+        Returns:
+            matplotlib.figure.Figure or list of figures
+        """
+        pass
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Plot images to TensorBoard."""
+        if epoch % self.freq != 0:
+            return
+
+        # Get validation data and outputs from kwargs
+        val_data = kwargs.get('val_data')
+        val_outputs = kwargs.get('val_outputs')
+
+        if val_data is None or val_outputs is None:
+            return
+
+        # Create figure(s)
+        figures = self.create_figure(val_data, val_outputs, epoch, **kwargs)
+
+        # Log to TensorBoard
+        if isinstance(figures, (list, tuple)):
+            # Multiple figures
+            for i, fig in enumerate(figures):
+                idx = self.batch_indices[i] if i < len(self.batch_indices) else i
+                self.writer.add_figure(
+                    f'{self.tag}_batch_{idx}',
+                    fig,
+                    global_step=epoch,
+                    close=True
+                )
+        else:
+            # Single figure
+            self.writer.add_figure(
+                self.tag,
+                figures,
+                global_step=epoch,
+                close=True
+            )
+
+        self.writer.flush()
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close TensorBoard writer."""
+        if self.writer:
+            self.writer.close()
+
+
+# ============================================================================
+# EXAMPLE IMAGE PLOTTERS
+# ============================================================================
+
+class SimpleImagePlotter(BaseImagePlotter):
+    """
+    Simple image plotter for input/output comparison.
+
+    Shows input, ground truth, and prediction side by side.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+        tag: Tag for the image in TensorBoard
+        freq: Frequency of plotting (every N epochs)
+        batch_indices: Which batch indices to visualize
+        cmap: Colormap for matplotlib (default: 'viridis')
+
+    Example:
+        plotter = SimpleImagePlotter(
+            'logs/images',
+            tag='validation',
+            batch_indices=[0, 1, 2]
+        )
+    """
+
+    def __init__(
+            self,
+            log_dir: str,
+            tag: str = 'validation',
+            freq: int = 1,
+            batch_indices: Optional[List[int]] = None,
+            cmap: str = 'viridis'
+    ):
+        super().__init__(log_dir, tag, freq, batch_indices)
+        self.cmap = cmap
+
+    def create_figure(self, val_data, val_outputs, epoch, **kwargs):
+        """Create side-by-side comparison figure."""
+        # Extract input and target
+        val_input = val_data[0]  # (B, C, H, W) or (B, T, C, H, W)
+        val_target = val_data[1]
+
+        # Create figures for each batch index
+        figures = []
+        for idx in self.batch_indices:
+            if idx >= val_input.shape[0]:
+                break
+
+            fig, axes = plt.subplots(1, 3, figsize=(15, 5))
+
+            # Get data for this batch index
+            input_img = self._prepare_image(val_input[idx])
+            target_img = self._prepare_image(val_target[idx])
+            pred_img = self._prepare_image(val_outputs[idx])
+
+            # Plot
+            axes[0].imshow(input_img, cmap=self.cmap)
+            axes[0].set_title('Input')
+            axes[0].axis('off')
+
+            axes[1].imshow(target_img, cmap=self.cmap)
+            axes[1].set_title('Ground Truth')
+            axes[1].axis('off')
+
+            axes[2].imshow(pred_img, cmap=self.cmap)
+            axes[2].set_title('Prediction')
+            axes[2].axis('off')
+
+            plt.tight_layout()
+            figures.append(fig)
+
+        return figures if len(figures) > 1 else figures[0]
+
+    def _prepare_image(self, tensor):
+        """
+        Prepare tensor for visualization.
+
+        Handles:
+        - Time series: (T, C, H, W) -> take middle frame
+        - Multi-channel: (C, H, W) -> take first channel
+        - Single channel: (1, H, W) -> squeeze
+        """
+        # Move to CPU and detach
+        img = tensor.detach().cpu()
+
+        # Handle time series
+        if img.ndim == 4:  # (T, C, H, W)
+            img = img[img.shape[0] // 2]  # Take middle frame
+
+        # Handle channels
+        if img.ndim == 3:  # (C, H, W)
+            if img.shape[0] == 1:
+                img = img[0]  # Single channel
+            elif img.shape[0] == 3:
+                img = img.permute(1, 2, 0)  # RGB
+            else:
+                img = img[0]  # Take first channel
+
+        return img.numpy()
+
+
+class DifferencePlotter(BaseImagePlotter):
+    """
+    Plot input, prediction, and difference (error) map.
+
+    Useful for regression tasks to visualize prediction errors.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+        tag: Tag for the image in TensorBoard
+        freq: Frequency of plotting (every N epochs)
+        batch_indices: Which batch indices to visualize
+
+    Example:
+        plotter = DifferencePlotter('logs/images', tag='errors')
+    """
+
+    def __init__(
+            self,
+            log_dir: str,
+            tag: str = 'difference',
+            freq: int = 1,
+            batch_indices: Optional[List[int]] = None
+    ):
+        super().__init__(log_dir, tag, freq, batch_indices)
+
+    def create_figure(self, val_data, val_outputs, epoch, **kwargs):
+        """Create figure with input, prediction, and difference."""
+        val_input = val_data[0]
+        val_target = val_data[1]
+
+        figures = []
+        for idx in self.batch_indices:
+            if idx >= val_input.shape[0]:
+                break
+
+            fig, axes = plt.subplots(1, 4, figsize=(20, 5))
+
+            # Get data
+            input_img = self._prepare_image(val_input[idx])
+            target_img = self._prepare_image(val_target[idx])
+            pred_img = self._prepare_image(val_outputs[idx])
+            diff_img = target_img - pred_img
+
+            # Plot
+            axes[0].imshow(input_img, cmap='gray')
+            axes[0].set_title('Input')
+            axes[0].axis('off')
+
+            axes[1].imshow(target_img, cmap='gray')
+            axes[1].set_title('Target')
+            axes[1].axis('off')
+
+            axes[2].imshow(pred_img, cmap='gray')
+            axes[2].set_title('Prediction')
+            axes[2].axis('off')
+
+            im = axes[3].imshow(diff_img, cmap='RdBu_r', vmin=-diff_img.std(), vmax=diff_img.std())
+            axes[3].set_title('Difference (Target - Pred)')
+            axes[3].axis('off')
+            plt.colorbar(im, ax=axes[3])
+
+            plt.tight_layout()
+            figures.append(fig)
+
+        return figures if len(figures) > 1 else figures[0]
+
+    def _prepare_image(self, tensor):
+        """Prepare tensor for visualization."""
+        img = tensor.detach().cpu()
+
+        # Handle time series
+        if img.ndim == 4:
+            img = img[img.shape[0] // 2]
+
+        # Handle channels
+        if img.ndim == 3:
+            if img.shape[0] > 1:
+                img = img[0]  # Take first channel
+            else:
+                img = img[0]
+
+        return img.numpy()

kito/callbacks/tensorboard_callbacks.py

@@ -0,0 +1,132 @@
+from typing import Optional, Callable
+
+from kito.callbacks.callback_base import Callback
+from torch.utils.tensorboard import SummaryWriter
+
+
+class TensorBoardScalars(Callback):
+    """
+    Log scalar metrics to TensorBoard.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+
+    Example:
+        tb_scalars = TensorBoardScalars('logs/tensorboard')
+    """
+
+    def __init__(self, log_dir: str):
+        self.log_dir = log_dir
+        self.writer = None
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Initialize TensorBoard writer."""
+        self.writer = SummaryWriter(self.log_dir)
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Log scalars to TensorBoard."""
+        if logs is None:
+            return
+
+        for key, value in logs.items():
+            self.writer.add_scalar(key, value, epoch)
+
+        self.writer.flush()
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close TensorBoard writer."""
+        if self.writer:
+            self.writer.close()
+
+
+class TensorBoardHistograms(Callback):
+    """
+    Log model parameter histograms to TensorBoard.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+        freq: Frequency of histogram logging (every N epochs)
+
+    Example:
+        tb_histograms = TensorBoardHistograms('logs/tensorboard', freq=5)
+    """
+
+    def __init__(self, log_dir: str, freq: int = 1):
+        self.log_dir = log_dir
+        self.freq = freq
+        self.writer = None
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Initialize TensorBoard writer."""
+        self.writer = SummaryWriter(self.log_dir)
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Log parameter histograms."""
+        if epoch % self.freq != 0:
+            return
+
+        # Handle DDP
+        model_to_log = model.module if hasattr(model, 'module') else model
+
+        for name, param in model_to_log.named_parameters():
+            self.writer.add_histogram(name, param, epoch)
+
+        self.writer.flush()
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close TensorBoard writer."""
+        if self.writer:
+            self.writer.close()
+
+
+class TensorBoardGraph(Callback):
+    """
+    Log model graph to TensorBoard.
+
+    Args:
+        log_dir: Directory for TensorBoard logs
+        input_to_model: Function that returns sample input for the model
+
+    Example:
+        def get_input():
+            return torch.randn(1, 3, 64, 64)
+
+        tb_graph = TensorBoardGraph('logs/tensorboard', input_to_model=get_input)
+    """
+
+    def __init__(
+            self,
+            log_dir: str,
+            input_to_model: Optional[Callable] = None
+    ):
+        self.log_dir = log_dir
+        self.input_to_model = input_to_model
+        self.writer = None
+        self.logged = False
+
+    def on_train_begin(self, engine, model, **kwargs):
+        """Initialize TensorBoard writer and log graph."""
+        if self.logged:
+            return
+
+        self.writer = SummaryWriter(self.log_dir)
+
+        # Get sample input
+        if self.input_to_model:
+            sample_input = self.input_to_model()
+        elif hasattr(engine, 'get_sample_input'):
+            sample_input = engine.get_sample_input()
+        else:
+            return  # Can't log graph without input
+
+        # Handle DDP
+        model_to_log = model.module if hasattr(model, 'module') else model
+
+        self.writer.add_graph(model_to_log, sample_input)
+        self.writer.flush()
+        self.logged = True
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close TensorBoard writer."""
+        if self.writer:
+            self.writer.close()

kito/callbacks/txt_logger.py

@@ -0,0 +1,57 @@
+import datetime
+import os
+
+from kito.callbacks.callback_base import Callback
+
+
+class TextLogger(Callback):
+    """
+    Log training metrics to a text file.
+
+    Args:
+        filename: Path to log file
+        append: Append to existing file or overwrite
+
+    Example:
+        text_logger = TextLogger('logs/training.log')
+    """
+
+    def __init__(
+            self,
+            filename: str,
+            append: bool = True
+    ):
+        self.filename = filename
+        self.append = append
+        self.file = 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 log file."""
+        mode = 'a' if self.append else 'w'
+        self.file = open(self.filename, mode)
+
+        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        self.file.write(f"{'=' * 60}")
+        self.file.write(f"Training started at {timestamp}")
+        self.file.write(f"{'=' * 60}")
+        self.file.flush()
+
+    def on_epoch_end(self, epoch, engine, model, logs=None, **kwargs):
+        """Write metrics to log file."""
+        if logs is None:
+            return
+
+        self.file.write(f"Epoch {epoch}: ")
+        metrics_str = ", ".join(f"{k}={v:.4f}" for k, v in logs.items())
+        self.file.write(metrics_str + "\n")
+        self.file.flush()
+
+    def on_train_end(self, engine, model, **kwargs):
+        """Close log file."""
+        if self.file:
+            timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            self.file.write(f"\nTraining ended at {timestamp}\n")
+            self.file.close()

kito/config/moduleconfig.py

@@ -0,0 +1,201 @@
+"""
+Base configuration system for KitoModule framework.
+
+Users can extend these base configs with their own custom parameters.
+"""
+from dataclasses import dataclass, field
+from typing import Tuple, List, Optional, Dict, Any
+
+
+@dataclass
+class PreprocessingStepConfig:
+    """
+    Configuration for a single preprocessing step.
+
+    Args:
+        type: Name of preprocessing class (e.g., 'detrend', 'standardization')
+        params: Dictionary of parameters to pass to preprocessing class
+
+    Example:
+        >>> step = PreprocessingStepConfig(
+        ...     type='standardization',
+        ...     params={'mean': 0.5, 'std': 0.2}
+        ... )
+    """
+    type: str
+    params: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class DataConfig:
+    """
+    Data loading and preprocessing configuration.
+
+    This config defines:
+    - What dataset to use (H5, memory, custom)
+    - Where data is located
+    - How to initialize the dataset (flexible args)
+    - Memory loading strategy
+    - How to split data (train/val/test)
+    - What preprocessing to apply
+    - DataLoader settings
+    """
+
+    # Dataset configuration
+    dataset_type: str = 'h5dataset'  # 'h5dataset', 'memdataset', or custom
+
+    # Simple path (backward compatible)
+    dataset_path: str = ''
+
+    # Flexible initialization args (for custom datasets)
+    # If provided, takes precedence over dataset_path
+    # Allows any constructor signature: Dataset(**dataset_init_args)
+    dataset_init_args: Dict[str, Any] = field(default_factory=dict)
+
+    # Memory management
+    load_into_memory: bool = False  # Load entire dataset into RAM for faster training
+
+    # Splitting ratios
+    train_ratio: float = 0.8
+    val_ratio: float = 0.1
+    # test_ratio is implicit: 1 - train_ratio - val_ratio
+
+    # Total samples to use (None = use all available)
+    total_samples: Optional[int] = None
+
+    # Preprocessing pipeline
+    # List of preprocessing steps applied in order
+    preprocessing: List[PreprocessingStepConfig] = field(default_factory=list)
+
+    # DataLoader settings
+    num_workers: int = 0
+    prefetch_factor: int = 2
+    pin_memory: bool = False
+    persistent_workers: bool = False
+
+
+@dataclass
+class TrainingConfig:
+    """Core training parameters required by KitoModule."""
+
+    # Essential training parameters
+    learning_rate: float
+    n_train_epochs: int
+    batch_size: int
+    train_mode: bool  # True for training, False for inference
+
+    # Verbosity (0=silent, 1=progress bar, 2=detailed)
+    train_verbosity_level: int = 2
+    val_verbosity_level: int = 2
+    test_verbosity_level: int = 2
+
+    # Distributed training
+    distributed_training: bool = False
+    master_gpu_id: int = 0  # Only used if not distributed
+
+    # Weight initialization
+    initialize_model_with_saved_weights: bool = False
+
+    # device type initialization
+    device_type: str = "cuda"  # "cuda", "mps", or "cpu"
+
+    def __post_init__(self):
+        """Validate device_type after initialization."""
+        valid_devices = {"cuda", "mps", "cpu"}
+
+        if self.device_type not in valid_devices:
+            raise ValueError(
+                f"Invalid device_type: '{self.device_type}'. "
+                f"Must be one of {valid_devices}."
+            )
+
+        # Normalize to lowercase (user-friendly)
+        self.device_type = self.device_type.lower()
+
+
+@dataclass
+class ModelConfig:
+    """Core model parameters required by KitoModule."""
+
+    # Data dimensions
+    input_data_size: Tuple[int, ...]  # Flexible shape
+
+    # Loss and optimization
+    loss: str = ""
+
+    # Callbacks and logging
+    log_to_tensorboard: bool = False
+    save_model_weights: bool = False
+    text_logging: bool = False
+    csv_logging: bool = False
+    train_codename: str = "experiment"
+
+    # Weights
+    weight_load_path: str = ""
+
+    # Inference
+    save_inference_to_disk: bool = False
+    inference_filename: str = ""
+
+    # TensorBoard visualization (optional)
+    tensorboard_img_id: str = "training_viz"
+    batch_idx_viz: List[int] = field(default_factory=lambda: [0])
+
+
+@dataclass
+class WorkDirConfig:
+    """Working directory configuration."""
+    work_directory: str
+
+
+@dataclass
+class CallbacksConfig:
+    """
+    Configuration for Kito's built-in default callbacks.
+
+    For custom callbacks, use instead:
+        callbacks = engine.get_default_callbacks()
+        callbacks.append(MyCustomCallback())
+        engine.fit(..., callbacks=callbacks)
+    """
+
+    # === CSV Logger ===
+    enable_csv_logger: bool = True
+
+    # === Text Logger ===
+    enable_text_logger: bool = True
+
+    # === Model Checkpoint ===
+    enable_model_checkpoint: bool = True
+    checkpoint_monitor: str = 'val_loss'
+    checkpoint_mode: str = 'min'  # 'min' or 'max'
+    checkpoint_save_best_only: bool = True
+    checkpoint_verbose: bool = False
+
+    # === TensorBoard ===
+    enable_tensorboard: bool = False  # Master switch
+    tensorboard_scalars: bool = True
+    tensorboard_histograms: bool = True
+    tensorboard_histogram_freq: int = 5
+    tensorboard_graph: bool = True
+    tensorboard_images: bool = False
+    tensorboard_image_freq: int = 1
+    tensorboard_batch_indices: List[int] = field(default_factory=lambda: [0])
+
+
+@dataclass
+class KitoModuleConfig:
+    """
+    Base configuration container for KitoModule.
+
+    Contains all configuration sections:
+    - training: Training parameters
+    - model: Model architecture and settings
+    - workdir: Output directories
+    - data: Dataset and preprocessing
+    """
+    training: TrainingConfig
+    model: ModelConfig
+    workdir: WorkDirConfig
+    data: DataConfig
+    callbacks: CallbacksConfig