Module neuralops.tracking

Experiment Tracking – pluggable backends (local, TensorBoard, WandB, MLflow) with a unified API and HuggingFace Trainer callback.

Sub-modules

neuralops.tracking.backends

Backend implementations for experiment tracking. Each backend handles the specifics of initialisation, metric/param/artifact logging, and finalisation.

neuralops.tracking.callbacks

HuggingFace Trainer callback that logs metrics and parameters to an ExperimentTracker.

neuralops.tracking.tracker

Unified ExperimentTracker – select a backend via string or TrackingBackend enum. Provides a single, consistent API: log_metric, log_params, …

Classes

class ExperimentTracker (backend: str = 'local',
experiment_name: str = 'experiment',
log_dir: str = './experiments/logs',
project: str | None = None,
config_dict: dict[str, typing.Any] | None = None,
tracking_uri: str | None = None)

Expand source code

class ExperimentTracker:
    """
    A facade that delegates to the chosen tracking backend.

    Usage:
        tracker = ExperimentTracker(backend="wandb", project="my_project", experiment_name="run1")
        tracker.log_metric("loss", 0.5, step=100)
        tracker.finalize({"final_accuracy": 0.93})
    """

    BACKEND_MAP = {
        LOCAL_BACKEND: LocalTracker,
        TENSORBOARD_BACKEND: TensorBoardTracker,
        WANDB_BACKEND: WandBTracker,
        MLFLOW_BACKEND: MLflowTracker,
    }

    def __init__(
        self,
        backend: str = LOCAL_BACKEND,
        experiment_name: str = "experiment",
        log_dir: str = "./experiments/logs",
        project: str | None = None,
        config_dict: dict[str, Any] | None = None,
        tracking_uri: str | None = None,
    ) -> None:
        """
        Args:
            backend: "local", "tensorboard", "wandb", or "mlflow".
            experiment_name: Name of the experiment / run.
            log_dir: Directory for local / tensorboard logs.
            project: Project name (for WandB / MLflow).
            config_dict: Hyperparameter dictionary (for WandB).
            tracking_uri: MLflow tracking URI.
        """
        if hasattr(backend, "value"):
            backend = str(backend.value)
        elif not isinstance(backend, str):
            backend = str(backend)
        backend = backend.lower()
        if backend not in self.BACKEND_MAP:
            raise ValueError(
                f"Unknown backend '{backend}'. Available: {list(self.BACKEND_MAP.keys())}"
            )

        self.backend_name = backend
        self.experiment_name = experiment_name
        self.log_dir = log_dir
        self.project = project
        self.config_dict = config_dict or {}
        self.tracking_uri = tracking_uri

        # Instantiate the backend
        if backend == LOCAL_BACKEND:
            self._tracker = LocalTracker(log_dir=log_dir, experiment_name=experiment_name)
        elif backend == TENSORBOARD_BACKEND:
            self._tracker = TensorBoardTracker(log_dir=log_dir, experiment_name=experiment_name)
        elif backend == WANDB_BACKEND:
            self._tracker = WandBTracker(
                project=project or "default",
                experiment_name=experiment_name,
                config_dict=self.config_dict,
            )
        elif backend == MLFLOW_BACKEND:
            self._tracker = MLflowTracker(
                experiment_name=experiment_name, tracking_uri=tracking_uri
            )
        else:
            raise NotImplementedError(f"Backend '{backend}' not implemented")

    def log_metric(self, name: str, value: float, step: int | None = None) -> None:
        """Log a single metric (scalar)."""
        self._tracker.log_metric(name, value, step)

    def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
        """Log multiple metrics at once."""
        for k, v in metrics.items():
            self._tracker.log_metric(k, v, step)

    def log_param(self, key: str, value: Any) -> None:
        """Log a single hyperparameter."""
        self._tracker.log_params({key: value})

    def log_params(self, params: dict[str, Any]) -> None:
        """Log multiple hyperparameters."""
        self._tracker.log_params(params)

    def log_artifact(self, local_path: str) -> None:
        """Save or upload a file (model checkpoint, plot, etc.)."""
        self._tracker.log_artifact(local_path)

    def finalize(self, result: dict[str, Any] | None = None) -> None:
        """
        Mark the run as completed and store final results.
        After calling this, the tracker should not be used again.
        """
        self._tracker.finalize(result)
        self._tracker.close()

    def close(self) -> None:
        """Close the backend (useful for cleaning up without final results)."""
        self._tracker.close()

    def get_callback(self):
        """Return a transformers-compatible TrainerCallback for experiment tracking."""
        from transformers import TrainerCallback

        class TrackingCallback(TrainerCallback):
            def __init__(self, tracker_instance):
                """Execute __init__ operation."""
                self.tracker = tracker_instance

            def on_log(self, args, state, control, logs=None, **kwargs):
                """Execute on_log operation."""
                if logs:
                    step = state.global_step if state else None
                    self.tracker.log_metrics(logs, step=step)

        return TrackingCallback(self)

A facade that delegates to the chosen tracking backend.

Usage -----= tracker = ExperimentTracker(backend="wandb", project="my_project", experiment_name="run1") tracker.log_metric("loss", 0.5, step=100) tracker.finalize({"final_accuracy": 0.93})

Args

-----=

backend

"local", "tensorboard", "wandb", or "mlflow".

experiment_name

Name of the experiment / run.

log_dir

Directory for local / tensorboard logs.

project

Project name (for WandB / MLflow).

config_dict

Hyperparameter dictionary (for WandB).

tracking_uri

MLflow tracking URI.

Class variables

var BACKEND_MAP

The type of the None singleton.

Methods

def close(self) ‑> None

Expand source code

def close(self) -> None:
    """Close the backend (useful for cleaning up without final results)."""
    self._tracker.close()

Close the backend (useful for cleaning up without final results).

def finalize(self, result: dict[str, typing.Any] | None = None) ‑> None

Expand source code

def finalize(self, result: dict[str, Any] | None = None) -> None:
    """
    Mark the run as completed and store final results.
    After calling this, the tracker should not be used again.
    """
    self._tracker.finalize(result)
    self._tracker.close()

Mark the run as completed and store final results. After calling this, the tracker should not be used again.

def get_callback(self)

Expand source code

def get_callback(self):
    """Return a transformers-compatible TrainerCallback for experiment tracking."""
    from transformers import TrainerCallback

    class TrackingCallback(TrainerCallback):
        def __init__(self, tracker_instance):
            """Execute __init__ operation."""
            self.tracker = tracker_instance

        def on_log(self, args, state, control, logs=None, **kwargs):
            """Execute on_log operation."""
            if logs:
                step = state.global_step if state else None
                self.tracker.log_metrics(logs, step=step)

    return TrackingCallback(self)

Return a transformers-compatible TrainerCallback for experiment tracking.

def log_artifact(self, local_path: str) ‑> None

Expand source code

def log_artifact(self, local_path: str) -> None:
    """Save or upload a file (model checkpoint, plot, etc.)."""
    self._tracker.log_artifact(local_path)

Save or upload a file (model checkpoint, plot, etc.).

def log_metric(self, name: str, value: float, step: int | None = None) ‑> None

Expand source code

def log_metric(self, name: str, value: float, step: int | None = None) -> None:
    """Log a single metric (scalar)."""
    self._tracker.log_metric(name, value, step)

Log a single metric (scalar).

def log_metrics(self, metrics: dict[str, float], step: int | None = None) ‑> None

Expand source code

def log_metrics(self, metrics: dict[str, float], step: int | None = None) -> None:
    """Log multiple metrics at once."""
    for k, v in metrics.items():
        self._tracker.log_metric(k, v, step)

Log multiple metrics at once.

def log_param(self, key: str, value: Any) ‑> None

Expand source code

def log_param(self, key: str, value: Any) -> None:
    """Log a single hyperparameter."""
    self._tracker.log_params({key: value})

Log a single hyperparameter.

def log_params(self, params: dict[str, typing.Any]) ‑> None

Expand source code

def log_params(self, params: dict[str, Any]) -> None:
    """Log multiple hyperparameters."""
    self._tracker.log_params(params)

Log multiple hyperparameters.

class TrackerCallback (tracker: ExperimentTracker)

Expand source code

class TrackerCallback(TrainerCallback):
    """Automatically log training/evaluation metrics to the tracker."""

    def __init__(self, tracker: ExperimentTracker) -> None:
        self.tracker = tracker

    def on_log(self, args, state, control, logs=None, **kwargs) -> None:
        """On log."""
        if logs:
            step = state.global_step
            for k, v in logs.items():
                if isinstance(v, (int, float)):
                    self.tracker.log_metric(k, v, step=step)

    def on_evaluate(self, args, state, control, logs=None, **kwargs) -> None:
        """On evaluate."""
        if logs:
            step = state.global_step
            for k, v in logs.items():
                if isinstance(v, (int, float)):
                    self.tracker.log_metric(f"eval_{k}", v, step=step)

    def on_train_end(self, args, state, control, **kwargs) -> None:
        """On train end."""
        self.tracker.finalize()

Automatically log training/evaluation metrics to the tracker.

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_evaluate(self, args, state, control, logs=None, **kwargs) ‑> None

Expand source code

def on_evaluate(self, args, state, control, logs=None, **kwargs) -> None:
    """On evaluate."""
    if logs:
        step = state.global_step
        for k, v in logs.items():
            if isinstance(v, (int, float)):
                self.tracker.log_metric(f"eval_{k}", v, step=step)

On evaluate.

def on_log(self, args, state, control, logs=None, **kwargs) ‑> None

Expand source code

def on_log(self, args, state, control, logs=None, **kwargs) -> None:
    """On log."""
    if logs:
        step = state.global_step
        for k, v in logs.items():
            if isinstance(v, (int, float)):
                self.tracker.log_metric(k, v, step=step)

On log.

def on_train_end(self, args, state, control, **kwargs) ‑> None

Expand source code

def on_train_end(self, args, state, control, **kwargs) -> None:
    """On train end."""
    self.tracker.finalize()

On train end.

class TrackingBackend (*values)

Expand source code

class TrackingBackend(Enum):
    LOCAL = "local"
    TENSORBOARD = "tensorboard"
    WANDB = "wandb"
    MLFLOW = "mlflow"

Create a collection of name/value pairs.

Example enumeration:

>>> class Color(Enum):
...     RED = 1
...     BLUE = 2
...     GREEN = 3

Access them by:

• attribute access:

Color.RED

• value lookup:

Color(1)

• name lookup:

Color['RED']

Enumerations can be iterated over, and know how many members they have:

>>> len(Color)
3

>>> list(Color)
[<Color.RED: 1>, <Color.BLUE: 2>, <Color.GREEN: 3>]

Methods can be added to enumerations, and members can have their own attributes – see the documentation for details.

Ancestors

• enum.Enum

Class variables

var LOCAL

The type of the None singleton.

var MLFLOW

The type of the None singleton.

var TENSORBOARD

The type of the None singleton.

var WANDB

The type of the None singleton.