ins-pricing 0.2.9__py3-none-any.whl → 0.3.1__py3-none-any.whl

ins_pricing/modelling/core/bayesopt/utils/torch_trainer_mixin.py

@@ -0,0 +1,591 @@
+"""PyTorch training mixin with resource management and training loops.
+
+This module provides the TorchTrainerMixin class which is used by
+PyTorch-based trainers (ResNet, FT, GNN) for:
+- Resource profiling and memory management
+- Batch size computation and optimization
+- DataLoader creation with DDP support
+- Generic training and validation loops with AMP
+- Early stopping and loss curve plotting
+"""
+
+from __future__ import annotations
+
+import copy
+import ctypes
+import gc
+import math
+import os
+import time
+from contextlib import nullcontext
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+import optuna
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.cuda.amp import autocast, GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+
+# Try to import plotting functions
+try:
+    import matplotlib
+    if os.name != "nt" and not os.environ.get("DISPLAY") and not os.environ.get("MPLBACKEND"):
+        matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+    _MPL_IMPORT_ERROR: Optional[BaseException] = None
+except Exception as exc:
+    matplotlib = None
+    plt = None
+    _MPL_IMPORT_ERROR = exc
+
+try:
+    from ....plotting.diagnostics import plot_loss_curve as plot_loss_curve_common
+except Exception:
+    try:
+        from ins_pricing.plotting.diagnostics import plot_loss_curve as plot_loss_curve_common
+    except Exception:
+        plot_loss_curve_common = None
+
+# Import from other utils modules
+from .constants import EPS, compute_batch_size, tweedie_loss, ensure_parent_dir
+from .distributed_utils import DistributedUtils
+
+
+def _plot_skip(label: str) -> None:
+    """Print message when plot is skipped due to missing matplotlib."""
+    if _MPL_IMPORT_ERROR is not None:
+        print(f"[Plot] Skip {label}: matplotlib unavailable ({_MPL_IMPORT_ERROR}).", flush=True)
+    else:
+        print(f"[Plot] Skip {label}: matplotlib unavailable.", flush=True)
+
+
+class TorchTrainerMixin:
+    """Shared helpers for PyTorch tabular trainers.
+
+    Provides resource profiling, memory management, batch size optimization,
+    and standardized training loops with mixed precision and DDP support.
+
+    This mixin is used by ResNetTrainer, FTTrainer, and GNNTrainer.
+    """
+
+    def _device_type(self) -> str:
+        """Get device type (cpu/cuda/mps)."""
+        return getattr(self, "device", torch.device("cpu")).type
+
+    def _resolve_resource_profile(self) -> str:
+        """Determine resource usage profile.
+
+        Returns:
+            One of: 'throughput', 'memory_saving', or 'auto'
+        """
+        profile = getattr(self, "resource_profile", None)
+        if not profile:
+            profile = os.environ.get("BAYESOPT_RESOURCE_PROFILE", "auto")
+        profile = str(profile).strip().lower()
+        if profile in {"cpu", "mps", "cuda"}:
+            profile = "auto"
+        if profile not in {"auto", "throughput", "memory_saving"}:
+            profile = "auto"
+        if profile == "auto" and self._device_type() == "cuda":
+            profile = "throughput"
+        return profile
+
+    def _log_resource_summary_once(self, profile: str) -> None:
+        """Log resource configuration summary once."""
+        if getattr(self, "_resource_summary_logged", False):
+            return
+        if dist.is_initialized() and not DistributedUtils.is_main_process():
+            return
+        self._resource_summary_logged = True
+        device = getattr(self, "device", torch.device("cpu"))
+        device_type = self._device_type()
+        cpu_count = os.cpu_count() or 1
+        cuda_count = torch.cuda.device_count() if torch.cuda.is_available() else 0
+        mps_available = bool(getattr(torch.backends, "mps", None) and torch.backends.mps.is_available())
+        ddp_enabled = bool(getattr(self, "is_ddp_enabled", False))
+        data_parallel = bool(getattr(self, "use_data_parallel", False))
+        print(
+            f">>> Resource summary: device={device}, device_type={device_type}, "
+            f"cpu_count={cpu_count}, cuda_count={cuda_count}, mps={mps_available}, "
+            f"ddp={ddp_enabled}, data_parallel={data_parallel}, profile={profile}"
+        )
+
+    def _available_system_memory(self) -> Optional[int]:
+        """Get available system RAM in bytes."""
+        if os.name == "nt":
+            class _MemStatus(ctypes.Structure):
+                _fields_ = [
+                    ("dwLength", ctypes.c_ulong),
+                    ("dwMemoryLoad", ctypes.c_ulong),
+                    ("ullTotalPhys", ctypes.c_ulonglong),
+                    ("ullAvailPhys", ctypes.c_ulonglong),
+                    ("ullTotalPageFile", ctypes.c_ulonglong),
+                    ("ullAvailPageFile", ctypes.c_ulonglong),
+                    ("ullTotalVirtual", ctypes.c_ulonglong),
+                    ("ullAvailVirtual", ctypes.c_ulonglong),
+                    ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
+                ]
+            status = _MemStatus()
+            status.dwLength = ctypes.sizeof(_MemStatus)
+            if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(status)):
+                return int(status.ullAvailPhys)
+            return None
+        try:
+            pages = os.sysconf("SC_AVPHYS_PAGES")
+            page_size = os.sysconf("SC_PAGE_SIZE")
+            return int(pages * page_size)
+        except Exception:
+            return None
+
+    def _available_cuda_memory(self) -> Optional[int]:
+        """Get available CUDA memory in bytes."""
+        if not torch.cuda.is_available():
+            return None
+        try:
+            free_mem, _total_mem = torch.cuda.mem_get_info()
+        except Exception:
+            return None
+        return int(free_mem)
+
+    def _estimate_sample_bytes(self, dataset) -> Optional[int]:
+        """Estimate memory per sample in bytes."""
+        try:
+            if len(dataset) == 0:
+                return None
+            sample = dataset[0]
+        except Exception:
+            return None
+
+        def _bytes(obj) -> int:
+            if obj is None:
+                return 0
+            if torch.is_tensor(obj):
+                return int(obj.element_size() * obj.nelement())
+            if isinstance(obj, np.ndarray):
+                return int(obj.nbytes)
+            if isinstance(obj, (list, tuple)):
+                return int(sum(_bytes(item) for item in obj))
+            if isinstance(obj, dict):
+                return int(sum(_bytes(item) for item in obj.values()))
+            return 0
+
+        sample_bytes = _bytes(sample)
+        return int(sample_bytes) if sample_bytes > 0 else None
+
+    def _cap_batch_size_by_memory(self, dataset, batch_size: int, profile: str) -> int:
+        """Cap batch size based on available memory."""
+        if batch_size <= 1:
+            return batch_size
+        sample_bytes = self._estimate_sample_bytes(dataset)
+        if sample_bytes is None:
+            return batch_size
+        device_type = self._device_type()
+        if device_type == "cuda":
+            available = self._available_cuda_memory()
+            if available is None:
+                return batch_size
+            if profile == "throughput":
+                budget_ratio = 0.8
+                overhead = 8.0
+            elif profile == "memory_saving":
+                budget_ratio = 0.5
+                overhead = 14.0
+            else:
+                budget_ratio = 0.6
+                overhead = 12.0
+        else:
+            available = self._available_system_memory()
+            if available is None:
+                return batch_size
+            if profile == "throughput":
+                budget_ratio = 0.4
+                overhead = 1.8
+            elif profile == "memory_saving":
+                budget_ratio = 0.25
+                overhead = 3.0
+            else:
+                budget_ratio = 0.3
+                overhead = 2.6
+        budget = int(available * budget_ratio)
+        per_sample = int(sample_bytes * overhead)
+        if per_sample <= 0:
+            return batch_size
+        max_batch = max(1, int(budget // per_sample))
+        if max_batch < batch_size:
+            print(
+                f">>> Memory cap: batch_size {batch_size} -> {max_batch} "
+                f"(per_sample~{sample_bytes}B, budget~{budget // (1024**2)}MB)"
+            )
+        return min(batch_size, max_batch)
+
+    def _resolve_num_workers(self, max_workers: int, profile: Optional[str] = None) -> int:
+        """Determine number of DataLoader workers."""
+        if os.name == 'nt':
+            return 0
+        if getattr(self, "is_ddp_enabled", False):
+            return 0
+        profile = profile or self._resolve_resource_profile()
+        if profile == "memory_saving":
+            return 0
+        worker_cap = min(int(max_workers), os.cpu_count() or 1)
+        if self._device_type() == "mps":
+            worker_cap = min(worker_cap, 2)
+        return worker_cap
+
+    def _build_dataloader(self,
+                          dataset,
+                          N: int,
+                          base_bs_gpu: tuple,
+                          base_bs_cpu: tuple,
+                          min_bs: int = 64,
+                          target_effective_cuda: int = 1024,
+                          target_effective_cpu: int = 512,
+                          large_threshold: int = 200_000,
+                          mid_threshold: int = 50_000):
+        """Build DataLoader with adaptive batch size and worker configuration.
+
+        Returns:
+            Tuple of (dataloader, accum_steps)
+        """
+        profile = self._resolve_resource_profile()
+        self._log_resource_summary_once(profile)
+        batch_size = compute_batch_size(
+            data_size=len(dataset),
+            learning_rate=self.learning_rate,
+            batch_num=self.batch_num,
+            minimum=min_bs
+        )
+        gpu_large, gpu_mid, gpu_small = base_bs_gpu
+        cpu_mid, cpu_small = base_bs_cpu
+
+        if self._device_type() == 'cuda':
+            device_count = torch.cuda.device_count()
+            if getattr(self, "is_ddp_enabled", False):
+                device_count = 1
+            if device_count > 1:
+                min_bs = min_bs * device_count
+                print(
+                    f">>> Multi-GPU detected: {device_count} devices. Adjusted min_bs to {min_bs}.")
+
+            if N > large_threshold:
+                base_bs = gpu_large * device_count
+            elif N > mid_threshold:
+                base_bs = gpu_mid * device_count
+            else:
+                base_bs = gpu_small * device_count
+        else:
+            base_bs = cpu_mid if N > mid_threshold else cpu_small
+
+        batch_size = compute_batch_size(
+            data_size=len(dataset),
+            learning_rate=self.learning_rate,
+            batch_num=self.batch_num,
+            minimum=min_bs
+        )
+        batch_size = min(batch_size, base_bs, N)
+        batch_size = self._cap_batch_size_by_memory(
+            dataset, batch_size, profile)
+
+        target_effective_bs = target_effective_cuda if self._device_type(
+        ) == 'cuda' else target_effective_cpu
+        if getattr(self, "is_ddp_enabled", False):
+            world_size = max(1, DistributedUtils.world_size())
+            target_effective_bs = max(1, target_effective_bs // world_size)
+
+        world_size = getattr(self, "world_size", 1) if getattr(
+            self, "is_ddp_enabled", False) else 1
+        samples_per_rank = math.ceil(
+            N / max(1, world_size)) if world_size > 1 else N
+        steps_per_epoch = max(
+            1, math.ceil(samples_per_rank / max(1, batch_size)))
+        desired_accum = max(1, target_effective_bs // max(1, batch_size))
+        accum_steps = max(1, min(desired_accum, steps_per_epoch))
+
+        workers = self._resolve_num_workers(8, profile=profile)
+        prefetch_factor = None
+        if workers > 0:
+            prefetch_factor = 4 if profile == "throughput" else 2
+        persistent = workers > 0 and profile != "memory_saving"
+        print(
+            f">>> DataLoader config: Batch Size={batch_size}, Accum Steps={accum_steps}, "
+            f"Workers={workers}, Prefetch={prefetch_factor or 'off'}, Profile={profile}")
+        sampler = None
+        use_distributed_sampler = bool(
+            dist.is_initialized() and getattr(self, "is_ddp_enabled", False)
+        )
+        if use_distributed_sampler:
+            sampler = DistributedSampler(dataset, shuffle=True)
+            shuffle = False
+        else:
+            shuffle = True
+
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            sampler=sampler,
+            num_workers=workers,
+            pin_memory=(self._device_type() == 'cuda'),
+            persistent_workers=persistent,
+            **({"prefetch_factor": prefetch_factor} if prefetch_factor is not None else {}),
+        )
+        return dataloader, accum_steps
+
+    def _build_val_dataloader(self, dataset, train_dataloader, accum_steps):
+        """Build validation DataLoader."""
+        profile = self._resolve_resource_profile()
+        val_bs = accum_steps * train_dataloader.batch_size
+        val_workers = self._resolve_num_workers(4, profile=profile)
+        prefetch_factor = None
+        if val_workers > 0:
+            prefetch_factor = 2
+        return DataLoader(
+            dataset,
+            batch_size=val_bs,
+            shuffle=False,
+            num_workers=val_workers,
+            pin_memory=(self._device_type() == 'cuda'),
+            persistent_workers=(val_workers > 0 and profile != "memory_saving"),
+            **({"prefetch_factor": prefetch_factor} if prefetch_factor is not None else {}),
+        )
+
+    def _compute_losses(self, y_pred, y_true, apply_softplus: bool = False):
+        """Compute per-sample losses based on task type."""
+        task = getattr(self, "task_type", "regression")
+        if task == 'classification':
+            loss_fn = nn.BCEWithLogitsLoss(reduction='none')
+            return loss_fn(y_pred, y_true).view(-1)
+        if apply_softplus:
+            y_pred = F.softplus(y_pred)
+        y_pred = torch.clamp(y_pred, min=1e-6)
+        power = getattr(self, "tw_power", 1.5)
+        return tweedie_loss(y_pred, y_true, p=power).view(-1)
+
+    def _compute_weighted_loss(self, y_pred, y_true, weights, apply_softplus: bool = False):
+        """Compute weighted loss."""
+        losses = self._compute_losses(
+            y_pred, y_true, apply_softplus=apply_softplus)
+        weighted_loss = (losses * weights.view(-1)).sum() / \
+            torch.clamp(weights.sum(), min=EPS)
+        return weighted_loss
+
+    def _early_stop_update(self, val_loss, best_loss, best_state, patience_counter, model,
+                           ignore_keys: Optional[List[str]] = None):
+        """Update early stopping state."""
+        if val_loss < best_loss:
+            ignore_keys = ignore_keys or []
+            base_module = model.module if hasattr(model, "module") else model
+            state_dict = {
+                k: (v.clone() if isinstance(v, torch.Tensor) else copy.deepcopy(v))
+                for k, v in base_module.state_dict().items()
+                if not any(k.startswith(ignore_key) for ignore_key in ignore_keys)
+            }
+            return val_loss, state_dict, 0, False
+        patience_counter += 1
+        should_stop = best_state is not None and patience_counter >= getattr(
+            self, "patience", 0)
+        return best_loss, best_state, patience_counter, should_stop
+
+    def _train_model(self,
+                     model,
+                     dataloader,
+                     accum_steps,
+                     optimizer,
+                     scaler,
+                     forward_fn,
+                     val_forward_fn=None,
+                     apply_softplus: bool = False,
+                     clip_fn=None,
+                     trial: Optional[optuna.trial.Trial] = None,
+                     loss_curve_path: Optional[str] = None):
+        """Generic training loop with AMP, DDP, and early stopping support.
+
+        Returns:
+            Tuple of (best_state_dict, history)
+        """
+        device_type = self._device_type()
+        best_loss = float('inf')
+        best_state = None
+        patience_counter = 0
+        stop_training = False
+        train_history: List[float] = []
+        val_history: List[float] = []
+
+        is_ddp_model = isinstance(model, DDP)
+        use_collectives = dist.is_initialized() and is_ddp_model
+
+        for epoch in range(1, getattr(self, "epochs", 1) + 1):
+            epoch_start_ts = time.time()
+            val_weighted_loss = None
+            if hasattr(self, 'dataloader_sampler') and self.dataloader_sampler is not None:
+                self.dataloader_sampler.set_epoch(epoch)
+
+            model.train()
+            optimizer.zero_grad()
+
+            epoch_loss_sum = None
+            epoch_weight_sum = None
+            for step, batch in enumerate(dataloader):
+                is_update_step = ((step + 1) % accum_steps == 0) or \
+                    ((step + 1) == len(dataloader))
+                sync_cm = model.no_sync if (
+                    is_ddp_model and not is_update_step) else nullcontext
+
+                with sync_cm():
+                    with autocast(enabled=(device_type == 'cuda')):
+                        y_pred, y_true, w = forward_fn(batch)
+                        weighted_loss = self._compute_weighted_loss(
+                            y_pred, y_true, w, apply_softplus=apply_softplus)
+                        loss_for_backward = weighted_loss / accum_steps
+
+                    batch_weight = torch.clamp(
+                        w.detach().sum(), min=EPS).to(dtype=torch.float32)
+                    loss_val = weighted_loss.detach().to(dtype=torch.float32)
+                    if epoch_loss_sum is None:
+                        epoch_loss_sum = torch.zeros(
+                            (), device=batch_weight.device, dtype=torch.float32)
+                        epoch_weight_sum = torch.zeros(
+                            (), device=batch_weight.device, dtype=torch.float32)
+                    epoch_loss_sum = epoch_loss_sum + loss_val * batch_weight
+                    epoch_weight_sum = epoch_weight_sum + batch_weight
+                    scaler.scale(loss_for_backward).backward()
+
+                if is_update_step:
+                    if clip_fn is not None:
+                        clip_fn()
+                    scaler.step(optimizer)
+                    scaler.update()
+                    optimizer.zero_grad()
+
+            if epoch_loss_sum is None or epoch_weight_sum is None:
+                train_epoch_loss = 0.0
+            else:
+                train_epoch_loss = (
+                    epoch_loss_sum / torch.clamp(epoch_weight_sum, min=EPS)
+                ).item()
+            train_history.append(float(train_epoch_loss))
+
+            if val_forward_fn is not None:
+                should_compute_val = (not dist.is_initialized()
+                                      or DistributedUtils.is_main_process())
+                val_device = getattr(self, "device", torch.device("cpu"))
+                if not isinstance(val_device, torch.device):
+                    val_device = torch.device(val_device)
+                loss_tensor_device = val_device if device_type == 'cuda' else torch.device(
+                    "cpu")
+                val_loss_tensor = torch.zeros(1, device=loss_tensor_device)
+
+                if should_compute_val:
+                    model.eval()
+                    with torch.no_grad(), autocast(enabled=(device_type == 'cuda')):
+                        val_result = val_forward_fn()
+                        if isinstance(val_result, tuple) and len(val_result) == 3:
+                            y_val_pred, y_val_true, w_val = val_result
+                            val_weighted_loss = self._compute_weighted_loss(
+                                y_val_pred, y_val_true, w_val, apply_softplus=apply_softplus)
+                        else:
+                            val_weighted_loss = val_result
+                    val_loss_tensor[0] = float(val_weighted_loss)
+
+                if use_collectives:
+                    dist.broadcast(val_loss_tensor, src=0)
+                val_weighted_loss = float(val_loss_tensor.item())
+
+                val_history.append(val_weighted_loss)
+
+                best_loss, best_state, patience_counter, stop_training = self._early_stop_update(
+                    val_weighted_loss, best_loss, best_state, patience_counter, model)
+
+                prune_flag = False
+                is_main_rank = DistributedUtils.is_main_process()
+                if trial is not None and is_main_rank:
+                    trial.report(val_weighted_loss, epoch)
+                    prune_flag = trial.should_prune()
+
+                if use_collectives:
+                    prune_device = getattr(self, "device", torch.device("cpu"))
+                    if not isinstance(prune_device, torch.device):
+                        prune_device = torch.device(prune_device)
+                    prune_tensor = torch.zeros(1, device=prune_device)
+                    if is_main_rank:
+                        prune_tensor.fill_(1 if prune_flag else 0)
+                    dist.broadcast(prune_tensor, src=0)
+                    prune_flag = bool(prune_tensor.item())
+
+                if prune_flag:
+                    raise optuna.TrialPruned()
+
+                if stop_training:
+                    break
+
+            should_log_epoch = (not dist.is_initialized()
+                                or DistributedUtils.is_main_process())
+            if should_log_epoch:
+                elapsed = int(time.time() - epoch_start_ts)
+                if val_weighted_loss is None:
+                    print(
+                        f"[Training] Epoch {epoch}/{getattr(self, 'epochs', 1)} "
+                        f"train_loss={float(train_epoch_loss):.6f} elapsed={elapsed}s",
+                        flush=True,
+                    )
+                else:
+                    print(
+                        f"[Training] Epoch {epoch}/{getattr(self, 'epochs', 1)} "
+                        f"train_loss={float(train_epoch_loss):.6f} "
+                        f"val_loss={float(val_weighted_loss):.6f} elapsed={elapsed}s",
+                        flush=True,
+                    )
+
+            if epoch % 10 == 0:
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                gc.collect()
+
+        history = {"train": train_history, "val": val_history}
+        self._plot_loss_curve(history, loss_curve_path)
+        return best_state, history
+
+    def _plot_loss_curve(self, history: Dict[str, List[float]], save_path: Optional[str]) -> None:
+        """Plot training and validation loss curves."""
+        if not save_path:
+            return
+        if dist.is_initialized() and not DistributedUtils.is_main_process():
+            return
+        train_hist = history.get("train", []) if history else []
+        val_hist = history.get("val", []) if history else []
+        if not train_hist and not val_hist:
+            return
+        if plot_loss_curve_common is not None:
+            plot_loss_curve_common(
+                history=history,
+                title="Loss vs. Epoch",
+                save_path=save_path,
+                show=False,
+            )
+        else:
+            if plt is None:
+                _plot_skip("loss curve")
+                return
+            ensure_parent_dir(save_path)
+            epochs = range(1, max(len(train_hist), len(val_hist)) + 1)
+            fig = plt.figure(figsize=(8, 4))
+            ax = fig.add_subplot(111)
+            if train_hist:
+                ax.plot(range(1, len(train_hist) + 1), train_hist,
+                        label='Train Loss', color='tab:blue')
+            if val_hist:
+                ax.plot(range(1, len(val_hist) + 1), val_hist,
+                        label='Validation Loss', color='tab:orange')
+            ax.set_xlabel('Epoch')
+            ax.set_ylabel('Weighted Loss')
+            ax.set_title('Loss vs. Epoch')
+            ax.grid(True, linestyle='--', alpha=0.3)
+            ax.legend()
+            plt.tight_layout()
+            plt.savefig(save_path, dpi=300)
+            plt.close(fig)
+        print(f"[Training] Loss curve saved to {save_path}")