torch-rechub 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

torch_rechub/basic/tracking.py

@@ -0,0 +1,198 @@
+"""Experiment tracking utilities for Torch-RecHub.
+
+This module exposes lightweight adapters for common visualization and
+experiment tracking tools, namely Weights & Biases (wandb), SwanLab, and
+TensorBoardX.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union
+
+
+class BaseLogger(ABC):
+    """Base interface for experiment tracking backends.
+
+    Methods
+    -------
+    log_metrics(metrics, step=None)
+        Record scalar metrics at a given step.
+    log_hyperparams(params)
+        Store hyperparameters and run configuration.
+    finish()
+        Flush pending logs and release resources.
+    """
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        """Log metrics to the tracking backend.
+
+        Parameters
+        ----------
+        metrics : dict of str to Any
+            Metric name-value pairs to record.
+        step : int, optional
+            Explicit global step or epoch index. When ``None``, the backend
+            uses its own default step handling.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        """Log experiment hyperparameters.
+
+        Parameters
+        ----------
+        params : dict of str to Any
+            Hyperparameters or configuration values to persist with the run.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def finish(self) -> None:
+        """Finalize logging and free any backend resources."""
+        raise NotImplementedError
+
+
+class WandbLogger(BaseLogger):
+    """Weights & Biases logger implementation.
+
+    Parameters
+    ----------
+    project : str
+        Name of the wandb project to log to.
+    name : str, optional
+        Display name for the run.
+    config : dict, optional
+        Initial hyperparameter configuration to record.
+    tags : list of str, optional
+        Optional tags for grouping runs.
+    notes : str, optional
+        Long-form notes shown in the run overview.
+    dir : str, optional
+        Local directory for wandb artifacts and cache.
+    **kwargs : dict
+        Additional keyword arguments forwarded to ``wandb.init``.
+
+    Raises
+    ------
+    ImportError
+        If ``wandb`` is not installed in the current environment.
+    """
+
+    def __init__(self, project: str, name: Optional[str] = None, config: Optional[Dict[str, Any]] = None, tags: Optional[List[str]] = None, notes: Optional[str] = None, dir: Optional[str] = None, **kwargs):
+        try:
+            import wandb
+            self._wandb = wandb
+        except ImportError:
+            raise ImportError("wandb is not installed. Install it with: pip install wandb")
+
+        self.run = self._wandb.init(project=project, name=name, config=config, tags=tags, notes=notes, dir=dir, **kwargs)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is not None:
+            self._wandb.log(metrics, step=step)
+        else:
+            self._wandb.log(metrics)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        if self.run is not None:
+            self.run.config.update(params)
+
+    def finish(self) -> None:
+        if self.run is not None:
+            self.run.finish()
+
+
+class SwanLabLogger(BaseLogger):
+    """SwanLab logger implementation.
+
+    Parameters
+    ----------
+    project : str, optional
+        Project identifier for grouping experiments.
+    experiment_name : str, optional
+        Display name for the experiment or run.
+    description : str, optional
+        Text description shown alongside the run.
+    config : dict, optional
+        Hyperparameters or configuration to log at startup.
+    logdir : str, optional
+        Directory where logs and artifacts are stored.
+    **kwargs : dict
+        Additional keyword arguments forwarded to ``swanlab.init``.
+
+    Raises
+    ------
+    ImportError
+        If ``swanlab`` is not installed in the current environment.
+    """
+
+    def __init__(self, project: Optional[str] = None, experiment_name: Optional[str] = None, description: Optional[str] = None, config: Optional[Dict[str, Any]] = None, logdir: Optional[str] = None, **kwargs):
+        try:
+            import swanlab
+            self._swanlab = swanlab
+        except ImportError:
+            raise ImportError("swanlab is not installed. Install it with: pip install swanlab")
+
+        self.run = self._swanlab.init(project=project, experiment_name=experiment_name, description=description, config=config, logdir=logdir, **kwargs)
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is not None:
+            self._swanlab.log(metrics, step=step)
+        else:
+            self._swanlab.log(metrics)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        if self.run is not None:
+            self.run.config.update(params)
+
+    def finish(self) -> None:
+        self._swanlab.finish()
+
+
+class TensorBoardXLogger(BaseLogger):
+    """TensorBoardX logger implementation.
+
+    Parameters
+    ----------
+    log_dir : str
+        Directory where event files will be written.
+    comment : str, default=""
+        Comment appended to the log directory name.
+    **kwargs : dict
+        Additional keyword arguments forwarded to
+        ``tensorboardX.SummaryWriter``.
+
+    Raises
+    ------
+    ImportError
+        If ``tensorboardX`` is not installed in the current environment.
+    """
+
+    def __init__(self, log_dir: str, comment: str = "", **kwargs):
+        try:
+            from tensorboardX import SummaryWriter
+            self._SummaryWriter = SummaryWriter
+        except ImportError:
+            raise ImportError("tensorboardX is not installed. Install it with: pip install tensorboardX")
+
+        self.writer = self._SummaryWriter(log_dir=log_dir, comment=comment, **kwargs)
+        self._step = 0
+
+    def log_metrics(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        if step is None:
+            step = self._step
+            self._step += 1
+
+        for key, value in metrics.items():
+            if value is not None:
+                if isinstance(value, (int, float)):
+                    self.writer.add_scalar(key, value, step)
+
+    def log_hyperparams(self, params: Dict[str, Any]) -> None:
+        hparam_str = "\n".join([f"{k}: {v}" for k, v in params.items()])
+        self.writer.add_text("hyperparameters", hparam_str, 0)
+
+    def finish(self) -> None:
+        if self.writer is not None:
+            self.writer.close()

torch_rechub/data/convert.py

@@ -0,0 +1,67 @@
+"""Utilities for converting array-like data structures into PyTorch tensors."""
+
+import numpy.typing as npt
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.types as pt
+import torch
+
+
+def pa_array_to_tensor(arr: pa.Array) -> torch.Tensor:
+    """
+    Convert a PyArrow array to a PyTorch tensor.
+
+    Parameters
+    ----------
+    arr : pa.Array
+        The given PyArrow array.
+
+    Returns
+    -------
+    torch.Tensor: The result PyTorch tensor.
+
+    Raises
+    ------
+    TypeError
+        if the array type or the value type (when nested) is unsupported.
+    ValueError
+        if the nested array is ragged (unequal lengths of each row).
+    """
+    if _is_supported_scalar(arr.type):
+        arr = pc.cast(arr, pa.float32())
+        return torch.from_numpy(_to_writable_numpy(arr))
+
+    if not _is_supported_list(arr.type):
+        raise TypeError(f"Unsupported array type: {arr.type}")
+
+    if not _is_supported_scalar(val_type := arr.type.value_type):
+        raise TypeError(f"Unsupported value type in the nested array: {val_type}")
+
+    if len(pc.unique(pc.list_value_length(arr))) > 1:
+        raise ValueError("Cannot convert the ragged nested array.")
+
+    arr = pc.cast(arr, pa.list_(pa.float32()))
+    np_arr = _to_writable_numpy(arr.values)  # type: ignore[attr-defined]
+
+    # For empty list-of-lists, define output shape as (0, 0); otherwise infer width.
+    return torch.from_numpy(np_arr.reshape(len(arr), -1 if len(arr) > 0 else 0))
+
+
+# helper functions
+
+
+def _is_supported_list(t: pa.DataType) -> bool:
+    """Check if the given PyArrow data type is a supported list."""
+    return pt.is_fixed_size_list(t) or pt.is_large_list(t) or pt.is_list(t)
+
+
+def _is_supported_scalar(t: pa.DataType) -> bool:
+    """Check if the given PyArrow data type is a supported scalar type."""
+    return pt.is_boolean(t) or pt.is_floating(t) or pt.is_integer(t) or pt.is_null(t)
+
+
+def _to_writable_numpy(arr: pa.Array) -> npt.NDArray:
+    """Dump a PyArrow array into a writable NumPy array."""
+    # Force the NumPy array to be writable. PyArrow's to_numpy() often returns a
+    # read-only view for zero-copy, which PyTorch's from_numpy() does not support.
+    return arr.to_numpy(writable=True, zero_copy_only=False)

torch_rechub/data/dataset.py

@@ -0,0 +1,120 @@
+"""Dataset implementations providing streaming, batch-wise data access for PyTorch."""
+
+import os
+import typing as ty
+
+import pyarrow.dataset as pd
+import torch
+from torch.utils.data import IterableDataset, get_worker_info
+
+from .convert import pa_array_to_tensor
+
+# Type for path to a file
+_FilePath = ty.Union[str, os.PathLike]
+
+# The default batch size when reading a Parquet dataset
+_DEFAULT_BATCH_SIZE = 1024
+
+
+class ParquetIterableDataset(IterableDataset):
+    """
+    IterableDataset that streams data from one or more Parquet files.
+
+    Parameters
+    ----------
+    file_paths : list[_FilePath]
+        Paths to Parquet files.
+    columns : list[str], optional
+        Column names to select. If ``None``, all columns are read.
+    batch_size : int, default DEFAULT_BATCH_SIZE
+        Number of rows per streamed batch.
+
+    Notes
+    -----
+    This dataset reads data lazily and never loads the entire Parquet dataset to memory.
+    The current worker receives a partition of ``file_paths`` and builds its own PyArrow
+    Dataset and Scanner. Iteration yields dictionaries mapping column names to PyTorch
+    tensors created via NumPy, one batch at a time.
+
+    Examples
+    --------
+    >>> ds = ParquetIterableDataset(
+    ...     ["/data/train1.parquet", "/data/train2.parquet"],
+    ...     columns=["x", "y", "label"],
+    ...     batch_size=1024,
+    ... )
+    >>> loader = DataLoader(ds, batch_size=None)
+    >>> # Now iterate over batches.
+    >>> for batch in loader:
+    ...     x, y, label = batch["x"], batch["y"], batch["label"]
+    ...     # Do some work.
+    ...     ...
+    """
+
+    def __init__(
+        self,
+        file_paths: ty.Sequence[_FilePath],
+        /,
+        columns: ty.Optional[ty.Sequence[str]] = None,
+        batch_size: int = _DEFAULT_BATCH_SIZE,
+    ) -> None:
+        """Initialize this instance."""
+        self._file_paths = tuple(map(str, file_paths))
+        self._columns = None if columns is None else tuple(columns)
+        self._batch_size = batch_size
+
+    def __iter__(self) -> ty.Iterator[dict[str, torch.Tensor]]:
+        """
+        Stream Parquet data as mapped PyTorch tensors.
+
+        Build a PyArrow Dataset from the current worker's assigned file partition, then
+        create a Scanner to lazily read batches of the selected columns. Each batch is
+        converted to a dict mapping column names to PyTorch tensors (via NumPy).
+
+        Returns
+        -------
+        Iterator[dict[str, torch.Tensor]]
+            An iterator that yields one converted batch at a time.
+        """
+        if not (partition := self._get_partition()):
+            return
+
+        # Build the dataset for the current worker.
+        ds = pd.dataset(partition, format="parquet")
+
+        # Create a scanner. This does not read data.
+        columns = None if self._columns is None else list(self._columns)
+        scanner = ds.scanner(columns=columns, batch_size=self._batch_size)
+
+        for batch in scanner.to_batches():
+            data_dict: dict[str, torch.Tensor] = {}
+            for name, array in zip(batch.column_names, batch.columns):
+                data_dict[name] = pa_array_to_tensor(array)
+            yield data_dict
+
+    # private interfaces
+
+    def _get_partition(self) -> tuple[str, ...]:
+        """
+        Get the partition of file paths for the current worker.
+
+        This method splits the full list of file paths into contiguous partitions with
+        a nearly equal size by the total number of workers and the current worker ID.
+
+        If running in the main process (i.e., no worker information is available), the
+        entire list of file paths is returned.
+
+        Returns
+        -------
+        tuple[str, ...]
+            The partition of file paths for the current worker.
+        """
+        if (info := get_worker_info()) is None:
+            return self._file_paths
+
+        n = len(self._file_paths)
+        per_worker = (n + info.num_workers - 1) // info.num_workers
+
+        start = info.id * per_worker
+        end = n if (end := start + per_worker) > n else end
+        return self._file_paths[start:end]

torch_rechub/trainers/ctr_trainer.py

@@ -43,6 +43,7 @@ class CTRTrainer(object):
         gpus=None,
         loss_mode=True,
         model_path="./",
+        model_logger=None,
     ):
         self.model = model  # for uniform weights save method in one gpu or multi gpu
         if gpus is None:
@@ -70,10 +71,13 @@ class CTRTrainer(object):
         self.model_path = model_path
         # Initialize regularization loss
         self.reg_loss_fn = RegularizationLoss(**regularization_params)
+        self.model_logger = model_logger
 
     def train_one_epoch(self, data_loader, log_interval=10):
         self.model.train()
         total_loss = 0
+        epoch_loss = 0
+        batch_count = 0
         tk0 = tqdm.tqdm(data_loader, desc="train", smoothing=0, mininterval=1.0)
         for i, (x_dict, y) in enumerate(tk0):
             x_dict = {k: v.to(self.device) for k, v in x_dict.items()}  # tensor to GPU
@@ -93,27 +97,62 @@ class CTRTrainer(object):
             loss.backward()
             self.optimizer.step()
             total_loss += loss.item()
+            epoch_loss += loss.item()
+            batch_count += 1
             if (i + 1) % log_interval == 0:
                 tk0.set_postfix(loss=total_loss / log_interval)
                 total_loss = 0
 
+        # Return average epoch loss
+        return epoch_loss / batch_count if batch_count > 0 else 0
+
     def fit(self, train_dataloader, val_dataloader=None):
+        for logger in self._iter_loggers():
+            logger.log_hyperparams({'n_epoch': self.n_epoch, 'learning_rate': self.optimizer.param_groups[0]['lr'], 'loss_mode': self.loss_mode})
+
         for epoch_i in range(self.n_epoch):
             print('epoch:', epoch_i)
-            self.train_one_epoch(train_dataloader)
+            train_loss = self.train_one_epoch(train_dataloader)
+
+            for logger in self._iter_loggers():
+                logger.log_metrics({'train/loss': train_loss, 'learning_rate': self.optimizer.param_groups[0]['lr']}, step=epoch_i)
+
             if self.scheduler is not None:
                 if epoch_i % self.scheduler.step_size == 0:
                     print("Current lr : {}".format(self.optimizer.state_dict()['param_groups'][0]['lr']))
                 self.scheduler.step()  # update lr in epoch level by scheduler
+
             if val_dataloader:
                 auc = self.evaluate(self.model, val_dataloader)
                 print('epoch:', epoch_i, 'validation: auc:', auc)
+
+                for logger in self._iter_loggers():
+                    logger.log_metrics({'val/auc': auc}, step=epoch_i)
+
                 if self.early_stopper.stop_training(auc, self.model.state_dict()):
                     print(f'validation: best auc: {self.early_stopper.best_auc}')
                     self.model.load_state_dict(self.early_stopper.best_weights)
                     break
+
         torch.save(self.model.state_dict(), os.path.join(self.model_path, "model.pth"))  # save best auc model
 
+        for logger in self._iter_loggers():
+            logger.finish()
+
+    def _iter_loggers(self):
+        """Return logger instances as a list.
+
+        Returns
+        -------
+        list
+            Active logger instances. Empty when ``model_logger`` is ``None``.
+        """
+        if self.model_logger is None:
+            return []
+        if isinstance(self.model_logger, (list, tuple)):
+            return list(self.model_logger)
+        return [self.model_logger]
+
     def evaluate(self, model, data_loader):
         model.eval()
         targets, predicts = list(), list()
@@ -189,3 +228,100 @@ class CTRTrainer(object):
 
         exporter = ONNXExporter(model, device=export_device)
         return exporter.export(output_path=output_path, dummy_input=dummy_input, batch_size=batch_size, seq_length=seq_length, opset_version=opset_version, dynamic_batch=dynamic_batch, verbose=verbose)
+
+    def visualization(self, input_data=None, batch_size=2, seq_length=10, depth=3, show_shapes=True, expand_nested=True, save_path=None, graph_name="model", device=None, dpi=300, **kwargs):
+        """Visualize the model's computation graph.
+
+        This method generates a visual representation of the model architecture,
+        showing layer connections, tensor shapes, and nested module structures.
+        It automatically extracts feature information from the model.
+
+        Parameters
+        ----------
+        input_data : dict, optional
+            Example input dict {feature_name: tensor}.
+            If not provided, dummy inputs will be generated automatically.
+        batch_size : int, default=2
+            Batch size for auto-generated dummy input.
+        seq_length : int, default=10
+            Sequence length for SequenceFeature.
+        depth : int, default=3
+            Visualization depth, higher values show more detail.
+            Set to -1 to show all layers.
+        show_shapes : bool, default=True
+            Whether to display tensor shapes.
+        expand_nested : bool, default=True
+            Whether to expand nested modules.
+        save_path : str, optional
+            Path to save the graph image (.pdf, .svg, .png).
+            If None, displays in Jupyter or opens system viewer.
+        graph_name : str, default="model"
+            Name for the graph.
+        device : str, optional
+            Device for model execution. If None, defaults to 'cpu'.
+        dpi : int, default=300
+            Resolution in dots per inch for output image.
+            Higher values produce sharper images suitable for papers.
+        **kwargs : dict
+            Additional arguments passed to torchview.draw_graph().
+
+        Returns
+        -------
+        ComputationGraph
+            A torchview ComputationGraph object.
+
+        Raises
+        ------
+        ImportError
+            If torchview or graphviz is not installed.
+
+        Notes
+        -----
+        Default Display Behavior:
+            When `save_path` is None (default):
+            - In Jupyter/IPython: automatically displays the graph inline
+            - In Python script: opens the graph with system default viewer
+
+        Examples
+        --------
+        >>> trainer = CTRTrainer(model, ...)
+        >>> trainer.fit(train_dl, val_dl)
+        >>>
+        >>> # Auto-display in Jupyter (no save_path needed)
+        >>> trainer.visualization(depth=4)
+        >>>
+        >>> # Save to high-DPI PNG for papers
+        >>> trainer.visualization(save_path="model.png", dpi=300)
+        """
+        from ..utils.visualization import TORCHVIEW_AVAILABLE, visualize_model
+
+        if not TORCHVIEW_AVAILABLE:
+            raise ImportError(
+                "Visualization requires torchview. "
+                "Install with: pip install torch-rechub[visualization]\n"
+                "Also ensure graphviz is installed on your system:\n"
+                "  - Ubuntu/Debian: sudo apt-get install graphviz\n"
+                "  - macOS: brew install graphviz\n"
+                "  - Windows: choco install graphviz"
+            )
+
+        # Handle DataParallel wrapped model
+        model = self.model.module if hasattr(self.model, 'module') else self.model
+
+        # Use provided device or default to 'cpu'
+        viz_device = device if device is not None else 'cpu'
+
+        return visualize_model(
+            model,
+            input_data=input_data,
+            batch_size=batch_size,
+            seq_length=seq_length,
+            depth=depth,
+            show_shapes=show_shapes,
+            expand_nested=expand_nested,
+            save_path=save_path,
+            graph_name=graph_name,
+            device=viz_device,
+            dpi=dpi,
+            **kwargs
+        )