pytorch-kito 0.2.0__py3-none-any.whl

kito/engine.py

@@ -0,0 +1,841 @@
+import datetime
+import os
+from pathlib import Path
+from typing import List, Optional
+
+import h5py
+import numpy as np
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+
+from kito.callbacks.callback_base import Callback, CallbackList
+from kito.callbacks.ddp_aware_callback import DDPAwareCallback
+from kito.config.moduleconfig import CallbacksConfig
+from kito.data.datapipeline import GenericDataPipeline
+from kito.module import KitoModule
+from kito.strategies.logger_strategy import DDPLogger, DefaultLogger
+from kito.strategies.progress_bar_strategy import (
+    StandardProgressBarHandler,
+    DDPProgressBarHandler
+)
+from kito.strategies.readiness_validator import ReadinessValidator
+from kito.utils.decorators import require_mode
+from kito.utils.gpu_utils import assign_device, get_available_devices
+
+
+class Engine:
+    """
+    Engine for training, validation, and inference.
+
+    HYBRID APPROACH:
+    - Auto-builds model if not built (user can also build explicitly)
+    - Auto-sets optimizer if not set (user can also set explicitly)
+
+    This provides:
+    1. Simplicity for beginners (just works)
+    2. Control for advanced users (explicit calls)
+    3. Flexibility for power users (custom args, modifications)
+
+    The Engine:
+    1. Manages device assignment and DDP
+    2. Iterates over batches and epochs
+    3. Calls module's single-batch methods (training_step, validation_step, prediction_step)
+    4. Manages callbacks (creates defaults, registers hooks)
+    5. Handles logging and progress bars
+
+    Args:
+        module: BaseModule instance
+        config: Configuration object (for device, DDP, callbacks, etc.)
+
+    Example (Simple - auto-build):
+        module = MyModel('MyModel', device, config)
+        engine = Engine(module, config)
+        engine.fit(train_loader, val_loader, max_epochs=100)  # Auto-builds and sets optimizer
+
+    Example (Advanced - explicit control):
+        module = MyModel('MyModel', device, config)
+        module.build(custom_layers=64)  # Custom build
+        module.summary()  # Inspect
+        module.associate_optimizer()  # Custom optimizer setup
+
+        engine = Engine(module, config)
+        engine.fit(train_loader, val_loader, max_epochs=100)  # Uses pre-built model
+    """
+
+    def __init__(self, module: KitoModule, config):
+        """
+        Initialize Engine.
+
+        Args:
+            module: KitoModule instance (can be built or not)
+            config: Configuration object
+        """
+        self.max_epochs = None
+        self.module = module
+        self.config = config
+
+        # Extract config values
+        self.distributed_training = config.training.distributed_training
+        self.work_directory = config.workdir.work_directory
+
+        # Logger
+        self.logger = DDPLogger() if self.distributed_training else DefaultLogger()
+
+        # Device and DDP setup
+        self.gpu_id = (
+            dist.get_rank()
+            if self.distributed_training
+            else config.training.master_gpu_id
+        )
+        self.driver_device = (
+            self.gpu_id == 0
+            if self.distributed_training
+            else True
+        )
+        device_type = config.training.device_type
+        self.device = assign_device(device_type, self.gpu_id)
+
+        # Log device info
+        available = get_available_devices()
+        self.logger.log_info(
+            f"Device configuration:\n"
+            f"  Requested: {device_type}\n"
+            f"  Assigned: {self.device}\n"
+            f"  Available: {available}"
+        )
+
+        # Assign to module
+        self.module._move_to_device(self.device)
+
+        # Progress bars
+        self.train_pbar = (
+            DDPProgressBarHandler()
+            if self.distributed_training
+            else StandardProgressBarHandler()
+        )
+        self.val_pbar = (
+            DDPProgressBarHandler()
+            if self.distributed_training
+            else StandardProgressBarHandler()
+        )
+        self.inference_pbar = StandardProgressBarHandler()
+
+        # Training state
+        self.current_epoch = 0
+        self.stop_training = False
+
+        # First batch flag for data shape checking
+        self._first_train_batch = True
+
+    # ========================================================================
+    # AUTO-SETUP (Hybrid approach)
+    # ========================================================================
+
+    def _ensure_model_ready_for_training(self):
+        """
+        Ensure model is ready for training.
+
+        Auto-builds and auto-sets optimizer if needed.
+        Logs when doing so for transparency.
+        """
+        # Auto-build if needed
+        if not self.module.is_built:
+            self.logger.log_info(
+                f"Model '{self.module.model_name}' not built. Building automatically..."
+            )
+            self.module.build()
+            self.module._move_to_device(self.device)  # Then move to device
+            self.logger.log_info("✓ Model built successfully.")
+
+        # Auto-setup optimizer if needed
+        if not self.module.is_optimizer_set:
+            self.logger.log_info(
+                f"Optimizer not set for '{self.module.model_name}'. Setting up automatically..."
+            )
+            self.module.associate_optimizer()
+            self.logger.log_info(
+                f"✓ Optimizer configured: {self.module.optimizer.__class__.__name__}"
+            )
+
+    def _ensure_model_ready_for_inference(self, weight_path: Optional[str] = None):
+        """
+        Ensure model is ready for inference.
+
+        Auto-builds if needed.
+        Does NOT auto-load weights (user must do this explicitly).
+        """
+        # Auto-build if needed
+        if not self.module.is_built:
+            self.logger.log_info(
+                f"Model '{self.module.model_name}' not built. Building automatically..."
+            )
+            self.module.build()
+            self.logger.log_info("✓ Model built successfully.")
+
+        # Check weights are loaded (don't auto-load, user must be explicit)
+        if not self.module.is_weights_loaded:
+            self.logger.log_warning(
+                f"⚠ Weights not loaded for '{self.module.model_name}'. "
+                "Call module.load_weights() or engine.load_weights() before inference."
+            )
+
+    # ========================================================================
+    # FIT - Training + Validation
+    # ========================================================================
+    @require_mode('train')
+    def fit(
+            self,
+            train_loader: DataLoader = None,
+            val_loader: DataLoader = None,
+            data_pipeline: GenericDataPipeline = None,
+            max_epochs: Optional[int] = None,
+            callbacks: Optional[List[Callback]] = None,
+    ):
+        """
+        Train the module.
+
+        HYBRID: Auto-builds model and sets optimizer if not already done.
+
+        Args:
+            train_loader: Training DataLoader
+            val_loader: Validation DataLoader
+            data_pipeline: GenericDataPipeline instance
+            max_epochs: Maximum epochs (None = use config value)
+            callbacks: List of callbacks (None = create smart defaults)
+
+        Example (Simple):
+            engine.fit(train_loader, val_loader, max_epochs=100)
+            # Auto-builds and sets optimizer
+
+        Example (Advanced):
+            module.build(custom_layers=64)
+            module.associate_optimizer()
+            engine.fit(train_loader, val_loader, max_epochs=100)
+            # Uses pre-built model
+        """
+        # ===== HYBRID: Auto-setup if needed =====
+        self._ensure_model_ready_for_training()
+
+        if data_pipeline is not None:
+            train_loader = data_pipeline.train_dataloader()
+            val_loader = data_pipeline.val_dataloader()
+
+        if train_loader is None and data_pipeline is None:
+            raise ValueError("Must provide either train_loader or data_pipeline")
+
+        # Validate readiness (should always pass after auto-setup)
+        ReadinessValidator.check_for_training(self.module)
+        ReadinessValidator.check_data_loaders(train_loader, val_loader)
+
+        # Get max_epochs
+        if max_epochs is None:
+            max_epochs = self.config.training.n_train_epochs
+        self.max_epochs = max_epochs
+
+        # Wrap model for DDP if needed
+        if self.distributed_training:
+            self._wrap_model_ddp()
+
+        # Setup callbacks
+        if callbacks is None:
+            callbacks = self._create_default_callbacks()
+
+        # Wrap callbacks for DDP
+        if self.distributed_training:
+            callbacks = [DDPAwareCallback(cb) for cb in callbacks]
+
+        callbacks = CallbackList(callbacks)
+
+        # Log training info
+        self._log_training_info(max_epochs, len(callbacks.callbacks))
+
+        # Reset first batch flag
+        self._first_train_batch = True
+
+        # ===== HOOK: on_train_begin =====
+        callbacks.on_train_begin(engine=self, model=self.module.model)
+
+        try:
+            for epoch in range(max_epochs):
+                self.current_epoch = epoch + 1
+
+                # ===== HOOK: on_epoch_begin =====
+                callbacks.on_epoch_begin(
+                    epoch=self.current_epoch,
+                    engine=self,
+                    model=self.module.model
+                )
+
+                # Train epoch
+                train_loss = self._train_epoch(
+                    train_loader,
+                    self.config.training.train_verbosity_level
+                )
+
+                # Validate epoch
+                val_loss, val_data, val_outputs = self._validate_epoch(
+                    val_loader,
+                    self.config.training.val_verbosity_level
+                )
+
+                # Prepare logs
+                logs = {
+                    'train_loss': train_loss,
+                    'val_loss': val_loss
+                }
+
+                # ===== HOOK: on_epoch_end =====
+                callbacks.on_epoch_end(
+                    epoch=self.current_epoch,
+                    engine=self,
+                    model=self.module.model,
+                    logs=logs,
+                    val_data=val_data,
+                    val_outputs=val_outputs
+                )
+
+                # Check early stopping
+                if self.stop_training:
+                    self.logger.log_info(f"\nStopping training at epoch {self.current_epoch}")
+                    break
+
+        except KeyboardInterrupt:
+            self.logger.log_info(f"\nTraining interrupted at epoch {self.current_epoch}")
+
+        finally:
+            # ===== HOOK: on_train_end =====
+            callbacks.on_train_end(engine=self, model=self.module.model)
+
+            # Cleanup DDP
+            if self.distributed_training:
+                dist.destroy_process_group()
+
+        self.logger.log_info(f"\nTraining of {self.module.model_name} completed.")
+
+    def _train_epoch(self, train_loader, verbosity_level):
+        """
+        Train for one epoch.
+
+        Iterates over batches and calls module.training_step(batch).
+        """
+        self.module.model.train()
+
+        # Set epoch for DDP sampler
+        if self.distributed_training:
+            train_loader.sampler.set_epoch(self.current_epoch)
+
+        # Init progress bar
+        self.train_pbar.init(
+            len(train_loader),
+            verbosity_level,
+            message= f"Epoch {self.current_epoch}/{self.max_epochs}"
+        )
+
+        # Accumulate loss
+        running_loss = 0.0
+
+        for batch_idx, batch in enumerate(train_loader):
+            # Check data shape on first batch
+            if self._first_train_batch:
+                self.module._check_data_shape()
+                self._first_train_batch = False
+
+            # ===== Call module's training_step =====
+            step_output = self.module.training_step(batch, self.train_pbar)
+
+            # Extract loss
+            loss = step_output['loss']
+            running_loss += loss.item()
+
+            # Update progress bar
+            self.train_pbar.step(
+                batch_idx + 1,
+                [("train_loss: ", float(f'{loss.item():.4f}'))]
+            )
+
+        # Return average loss
+        return running_loss / len(train_loader)
+
+    def _validate_epoch(self, val_loader, verbosity_level):
+        """
+        Validate for one epoch.
+
+        Iterates over batches and calls module.validation_step(batch).
+        """
+        self.module.model.eval()
+
+        # Init progress bar
+        self.val_pbar.init(
+            len(val_loader) + 1,
+            verbosity_level
+        )
+
+        # Accumulate metrics
+        running_loss = 0.0
+        last_inputs = None
+        last_targets = None
+        last_outputs = None
+
+        with torch.no_grad():
+            for batch_idx, batch in enumerate(val_loader):
+                # ===== Call module's validation_step =====
+                step_output = self.module.validation_step(batch, self.val_pbar)
+
+                # Extract metrics
+                loss = step_output['loss']
+                running_loss += loss.item()
+
+                # Store last batch for callbacks
+                last_outputs = step_output.get('outputs')
+                last_targets = step_output.get('targets')
+                last_inputs = step_output.get('inputs')
+
+                # Update progress bar
+                self.val_pbar.step(
+                    batch_idx + 1,
+                    [("val_loss: ", float(f'{loss.item():.4f}'))]
+                )
+
+        # Average loss
+        avg_loss = running_loss / len(val_loader)
+
+        # Package validation data for callbacks
+        val_data = (last_inputs, last_targets)
+
+        return avg_loss, val_data, last_outputs
+
+    # ========================================================================
+    # PREDICT - Inference
+    # ========================================================================
+    @require_mode('inference')
+    def predict(
+            self,
+            test_loader,
+            save_to_disk: bool = False,
+            output_path: Optional[str] = None,
+    ):
+        """
+        Run inference on test data.
+
+        HYBRID: Auto-builds model if not already built.
+        Does NOT auto-load weights (user must load explicitly).
+
+        Args:
+            test_loader: Test DataLoader
+            save_to_disk: Save predictions to HDF5 file
+            output_path: Path to save predictions (if save_to_disk=True)
+
+        Returns:
+            numpy.ndarray: Predictions (if save_to_disk=False)
+            None: If save_to_disk=True
+
+        Example:
+            # Load weights first (explicit)
+            module.load_weights('weights/best.pt')
+
+            # Predict (auto-builds if needed)
+            predictions = engine.predict(test_loader)
+        """
+        # ===== HYBRID: Auto-setup if needed =====
+        self._ensure_model_ready_for_inference()
+
+        # Validate (will warn if weights not loaded)
+        ReadinessValidator.check_for_inference(self.module)
+        ReadinessValidator.check_data_loaders(test_loader=test_loader)
+
+        if save_to_disk and output_path is None:
+            raise ValueError("output_path required when save_to_disk=True")
+
+        # Log inference info
+        self._log_inference_info(len(test_loader), save_to_disk, output_path)
+
+        # Prepare storage
+        storage = self._prepare_prediction_storage(
+            test_loader,
+            save_to_disk,
+            output_path
+        )
+
+        # Run inference
+        self.module.model.eval()
+
+        self.inference_pbar.init(
+            len(test_loader),
+            self.config.training.test_verbosity_level,
+            message='Evaluating model in inference mode'
+        )
+
+        with torch.no_grad():
+            for batch_idx, batch in enumerate(test_loader):
+                # ===== Call module's prediction_step =====
+                outputs = self.module.prediction_step(batch, self.inference_pbar)
+
+                # Store predictions
+                self._store_predictions(
+                    storage,
+                    batch_idx,
+                    outputs,
+                    test_loader.batch_size,
+                    len(test_loader.dataset)
+                )
+
+                # Update progress bar
+                self.inference_pbar.step(batch_idx + 1, values=None)
+
+        # Finalize storage
+        result = self._finalize_prediction_storage(storage, save_to_disk, output_path)
+
+        self.logger.log_info(f"\nInference of {self.module.model_name} completed.")
+
+        return result
+
+    def _prepare_prediction_storage(self, test_loader, save_to_disk, output_path):
+        """Prepare storage for predictions."""
+        # Infer output shape from first batch
+        sample_batch = next(iter(test_loader))
+
+        with torch.no_grad():
+            sample_output = self.module.prediction_step(sample_batch, None)
+
+        # Determine shapes
+        batch_size = sample_output.shape[0]
+        output_shape = sample_output.shape[1:]
+        total_samples = len(test_loader.dataset)
+        full_shape = (total_samples,) + output_shape
+
+        # Use module's standard_data_shape if set, otherwise infer
+        if self.module.standard_data_shape is not None:
+            output_shape = self.module.standard_data_shape
+            full_shape = (total_samples,) + output_shape
+
+        if save_to_disk:
+            # Validate output path
+            self._check_inference_save_path_valid(output_path)
+
+            # Create HDF5 file
+            h5_file = h5py.File(output_path, 'w')
+            h5_dataset = h5_file.create_dataset(
+                'predictions',
+                shape=full_shape,
+                dtype='float32',
+                chunks=(batch_size,) + output_shape
+            )
+
+            self.logger.log_info(f"Created HDF5 dataset '{output_path}' to store predictions.")
+
+            return {'type': 'disk', 'file': h5_file, 'dataset': h5_dataset}
+        else:
+            # Allocate memory
+            predictions = np.zeros(full_shape, dtype=np.float32)
+            self.logger.log_info(f"Created tensor in memory to store predictions.")
+            return {'type': 'memory', 'data': predictions}
+
+    def _store_predictions(self, storage, batch_idx, outputs, batch_size, dataset_len):
+        """Store predictions for a batch."""
+        batch_data = outputs.cpu().detach().numpy()
+
+        start_idx = batch_idx * batch_size
+        end_idx = min((batch_idx + 1) * batch_size, dataset_len)
+
+        if storage['type'] == 'disk':
+            storage['dataset'][start_idx:end_idx] = batch_data
+        else:
+            storage['data'][start_idx:end_idx] = batch_data
+
+    def _finalize_prediction_storage(self, storage, save_to_disk, output_path):
+        """Finalize storage and return result."""
+        if storage['type'] == 'disk':
+            storage['file'].close()
+            self.logger.log_info(f"Saved predictions to '{output_path}'")
+            return None
+        else:
+            return storage['data']
+
+    def _check_inference_save_path_valid(self, inference_filename):
+        """Validate inference output path."""
+        file_name, file_extension = os.path.splitext(inference_filename)
+
+        if os.path.isdir(inference_filename):
+            raise IsADirectoryError(
+                f"ERROR: '{os.path.abspath(inference_filename)}' is a directory."
+            )
+
+        if file_extension != '.h5':
+            raise ValueError(
+                f"ERROR: '{os.path.abspath(inference_filename)}' must have .h5 extension."
+            )
+
+        if os.path.exists(inference_filename):
+            raise FileExistsError(
+                f"ERROR: '{os.path.abspath(inference_filename)}' already exists."
+            )
+
+    # ========================================================================
+    # DDP
+    # ========================================================================
+
+    def _wrap_model_ddp(self):
+        """Wrap model with DistributedDataParallel."""
+        from torch.nn.parallel import DistributedDataParallel
+
+        if not isinstance(self.module.model, DistributedDataParallel):
+            self.module.model = DistributedDataParallel(
+                self.module.model,
+                device_ids=[self.gpu_id]
+            )
+            # self.module.model.to(self.device)
+            self.module._move_to_device(self.device)
+            self.logger.log_info("Model wrapped in DistributedDataParallel.")
+
+    # ========================================================================
+    # DEFAULT CALLBACKS
+    # ========================================================================
+
+    '''def _create_default_callbacks(self):
+        """Create smart default callbacks based on config."""
+          # vedi se tenerlo cosi' + forse fai interfacce in __init__.py
+        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, TensorBoardHistograms, TensorBoardGraph
+
+        from kito.callbacks.tensorboard_callback_images import SimpleImagePlotter
+
+        callbacks = []
+
+        # Setup paths
+        timestamp = datetime.datetime.now().strftime("%d%b%Y-%H%M%S")
+        work_dir = Path(os.path.expandvars(self.work_directory))
+        model_name = self.module.model_name
+        train_codename = self.config.model.train_codename
+
+        # === CSV Logger (always) ===
+        csv_dir = work_dir / "logs" / "csv"
+        csv_dir.mkdir(parents=True, exist_ok=True)
+        csv_path = csv_dir / f"{model_name}_{timestamp}_{train_codename}.csv"
+        callbacks.append(CSVLogger(str(csv_path)))
+
+        # === Text Logger (always) ===
+        log_dir = work_dir / "logs" / "text"
+        log_dir.mkdir(parents=True, exist_ok=True)
+        log_path = log_dir / f"{model_name}_{timestamp}_{train_codename}.log"
+        callbacks.append(TextLogger(str(log_path)))
+
+        # === Model Checkpoint (if configured) ===
+        if self.config.model.save_model_weights:
+            weight_dir = work_dir / "weights" / model_name
+            weight_dir.mkdir(parents=True, exist_ok=True)
+            weight_path = weight_dir / f"best_{model_name}_{timestamp}_{train_codename}.pt"
+
+            callbacks.append(
+                ModelCheckpoint(
+                    filepath=str(weight_path),
+                    monitor='val_loss',
+                    save_best_only=True,
+                    mode='min',
+                    verbose=True
+                )
+            )
+
+        # === TensorBoard (if configured) ===
+        if self.config.model.log_to_tensorboard:
+            tb_dir = work_dir / "logs" / "tensorboard" / model_name / timestamp / train_codename
+            tb_dir.mkdir(parents=True, exist_ok=True)
+
+            # Scalars
+            callbacks.append(TensorBoardScalars(str(tb_dir)))
+
+            # Histograms
+            callbacks.append(TensorBoardHistograms(str(tb_dir), freq=5))
+
+            # Model graph
+            callbacks.append(
+                TensorBoardGraph(
+                    str(tb_dir),
+                    input_to_model=lambda: self.module.get_sample_input()
+                )
+            )
+
+            # Image plotting (if configured)
+            if hasattr(self.config.model, 'tensorboard_img_id'):
+                img_dir = tb_dir / 'images'
+                callbacks.append(
+                    SimpleImagePlotter(
+                        log_dir=str(img_dir),
+                        tag=self.config.model.tensorboard_img_id,
+                        freq=1,
+                        batch_indices=getattr(self.config.model, 'batch_idx_viz', [0])
+                    )
+                )
+
+        return callbacks'''
+
+    def _create_default_callbacks(self):
+        """Create smart default callbacks based on config."""
+        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,
+            TensorBoardHistograms,
+            TensorBoardGraph
+        )
+        from kito.callbacks.tensorboard_callback_images import SimpleImagePlotter
+
+        callbacks = []
+
+        # Get callbacks config (with defaults if not provided)
+        cb_config = getattr(self.config, 'callbacks', CallbacksConfig())
+
+        # Setup paths
+        timestamp = datetime.datetime.now().strftime("%d%b%Y-%H%M%S")
+        work_dir = Path(os.path.expandvars(self.work_directory))
+        model_name = self.module.model_name
+        train_codename = self.config.model.train_codename
+
+        # === CSV Logger ===
+        if cb_config.enable_csv_logger:
+            csv_dir = work_dir / "logs" / "csv"
+            csv_dir.mkdir(parents=True, exist_ok=True)
+            csv_path = csv_dir / f"{model_name}_{timestamp}_{train_codename}.csv"
+            callbacks.append(CSVLogger(str(csv_path)))
+
+        # === Text Logger ===
+        if cb_config.enable_text_logger:
+            log_dir = work_dir / "logs" / "text"
+            log_dir.mkdir(parents=True, exist_ok=True)
+            log_path = log_dir / f"{model_name}_{timestamp}_{train_codename}.log"
+            callbacks.append(TextLogger(str(log_path)))
+
+        # === Model Checkpoint ===
+        if cb_config.enable_model_checkpoint and self.config.model.save_model_weights:  # should leave only one of the two
+            weight_dir = work_dir / "weights" / model_name
+            weight_dir.mkdir(parents=True, exist_ok=True)
+            weight_path = weight_dir / f"best_{model_name}_{timestamp}_{train_codename}.pt"
+
+            callbacks.append(
+                ModelCheckpoint(
+                    filepath=str(weight_path),
+                    monitor=cb_config.checkpoint_monitor,
+                    save_best_only=cb_config.checkpoint_save_best_only,
+                    mode=cb_config.checkpoint_mode,
+                    verbose=cb_config.checkpoint_verbose
+                )
+            )
+
+        # === TensorBoard ===
+        if cb_config.enable_tensorboard and self.config.model.log_to_tensorboard:  # should leave only one of the two
+            tb_dir = work_dir / "logs" / "tensorboard" / model_name / timestamp / train_codename
+            tb_dir.mkdir(parents=True, exist_ok=True)
+
+            # Scalars
+            if cb_config.tensorboard_scalars:
+                callbacks.append(TensorBoardScalars(str(tb_dir)))
+
+            # Histograms
+            if cb_config.tensorboard_histograms:
+                callbacks.append(
+                    TensorBoardHistograms(
+                        str(tb_dir),
+                        freq=cb_config.tensorboard_histogram_freq
+                    )
+                )
+
+            # Model graph
+            if cb_config.tensorboard_graph:
+                callbacks.append(
+                    TensorBoardGraph(
+                        str(tb_dir),
+                        input_to_model=lambda: self.module.get_sample_input()
+                    )
+                )
+
+            # Image plotting
+            if cb_config.tensorboard_images:
+                img_dir = tb_dir / 'images'
+                callbacks.append(
+                    SimpleImagePlotter(
+                        log_dir=str(img_dir),
+                        tag=getattr(self.config.model, 'tensorboard_img_id', 'images'),
+                        freq=cb_config.tensorboard_image_freq,
+                        batch_indices=cb_config.tensorboard_batch_indices
+                    )
+                )
+
+        return callbacks
+
+    def get_default_callbacks(self):
+        """
+        Get default callbacks configured from config.callbacks.
+
+        Useful when you want to extend defaults with custom callbacks.
+
+        Returns:
+            List of callback instances
+
+        Example:
+            >>> # Get defaults
+            >>> callbacks = engine.get_default_callbacks()
+            >>>
+            >>> # Modify or extend
+            >>> callbacks.append(MyCustomCallback(param=10))
+            >>>
+            >>> # Or modify existing
+            >>> for cb in callbacks:
+            ...     if isinstance(cb, ModelCheckpoint):
+            ...         cb.verbose = False
+            >>>
+            >>> # Use them
+            >>> engine.fit(train_loader, val_loader, callbacks=callbacks)
+        """
+        return self._create_default_callbacks()
+    # ========================================================================
+    # WEIGHT LOADING (Convenience method)
+    # ========================================================================
+
+    def load_weights(self, weight_path: str, strict: bool = True):
+        """
+        Load weights into module.
+
+        Convenience method that calls module.load_weights().
+
+        Args:
+            weight_path: Path to weight file
+            strict: Strict loading
+        """
+        self.module.load_weights(weight_path, strict)
+        self.logger.log_info(f"Loaded weights from '{weight_path}' successfully.")
+
+    # ========================================================================
+    # LOGGING
+    # ========================================================================
+
+    def _log_training_info(self, max_epochs, num_callbacks):
+        """Log training configuration."""
+        attrib_to_log = {
+            'model_name': self.module.model_name,
+            'optimizer': self.module.optimizer.__class__.__name__,
+            'batch_size': self.module.batch_size,
+            'n_train_epochs': max_epochs,
+            'learning_rate': self.module.learning_rate,
+            'work_directory': self.work_directory,
+            'save_model_weights': self.config.model.save_model_weights,
+            'distributed_training': self.distributed_training,
+            'callbacks': num_callbacks
+        }
+
+        self.logger.log_info(
+            'Model being used with the following parameters:\n\n' +
+            '\n '.join(f"{k} -> {v}" for k, v in attrib_to_log.items()) + '\n'
+        )
+
+    def _log_inference_info(self, num_batches, save_to_disk, output_path):
+        """Log inference configuration."""
+        self.logger.log_info(
+            f"\nRunning inference for {self.module.model_name}\n"
+            f"Batches: {num_batches}\n"
+            f"Device: {self.device}\n"
+            f"Save to disk: {save_to_disk}\n"
+            + (f"Output: {output_path}\n" if save_to_disk else "Output: In-memory\n")
+        )