Tracking API¶
torch_batteries.tracking
¶
Backend-neutral experiment tracking with optional W&B integration.
Public API¶
ExperimentTracker— interface implemented by tracking backends.Run— immutable run name, grouping, tags, description, and configuration.WandbTracker— optional Weights & Biases backend installed with thewandbextra.
ExperimentTracker
Source
¶
Bases: ABC
Abstract base class for experiment tracking backends.
Provides a unified interface for logging experiments, metrics, and artifacts to various tracking services (e.g. Weights & Biases).
The tracker is a standalone service that can be used independently or integrated with training via ExperimentTrackingCallback.
is_initialized
abstractmethod
property
¶
Check if the tracker has been initialized.
Returns:
| Name | Type | Description |
|---|---|---|
bool |
bool
|
True if initialized, False otherwise |
init(run) Source
abstractmethod
¶
Initialize the tracking session.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run
|
Run
|
Run configuration |
required |
log_metrics(metrics, step=None, prefix=None) Source
abstractmethod
¶
Log metrics to the tracker.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics
|
dict[str, float]
|
Dictionary of metric names to values |
required |
step
|
int | None
|
Optional step number (epoch, iteration, etc.) |
None
|
prefix
|
str | None
|
Optional prefix for metric names (e.g., "train/", "val/") |
None
|
finish(exit_code=0) Source
abstractmethod
¶
Finish tracking the experiment.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
exit_code
|
int
|
Exit code (0 for success, non-zero for failure) |
0
|
log_summary(summary) Source
abstractmethod
¶
Log summary of the experiment.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
summary
|
dict[str, Any]
|
Summary metrics/info |
required |
log_model(model, name='model', *, aliases=None, metadata=None) Source
abstractmethod
¶
Log a trained model artifact.
Trackers that support artifact logging (e.g. W&B) can override this.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
model
|
Module
|
Trained PyTorch model |
required |
name
|
str
|
Artifact base name |
'model'
|
aliases
|
list[str] | None
|
Optional artifact aliases (backend-specific) |
None
|
metadata
|
dict[str, Any] | None
|
Optional artifact metadata (backend-specific) |
None
|
Run
Source
dataclass
¶
A single run with a configuration.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
name
|
str | None
|
Optional run identifier |
None
|
group
|
str | None
|
Optional group name for organizing runs |
None
|
job_type
|
str | None
|
Optional type of job (e.g., "training", "evaluation") |
None
|
description
|
str | None
|
Optional description of the run |
None
|
config
|
dict[str, Any]
|
Run-specific configuration |
dict()
|
WandbTracker
Source
¶
Bases: ExperimentTracker
Weights & Biases experiment tracker implementation.
Example:
tracker = WandbTracker(project="your-wandb-project")
tracker.init(
run=Run(config={"lr": 0.001})
)
# During training
tracker.log_metrics({"train/loss": 0.5}, step=100)
tracker.finish()
run
property
¶
Get the tracked wandb run.
entity
property
¶
Get the wandb entity.
project
property
¶
Get the wandb project.
is_initialized
property
¶
Check if the tracker has been initialized.
Returns:
| Name | Type | Description |
|---|---|---|
bool |
bool
|
True if initialized, False otherwise |
run_id
property
¶
Get the current run ID.
run_url
property
¶
Get the wandb run URL.
__init__(project, entity=None) Source
¶
Initialize the wandb tracker.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
project
|
str
|
Wandb project name |
required |
entity
|
str | None
|
Optional wandb entity (username or team name) |
None
|
init(run) Source
¶
Initialize wandb tracking session.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
run
|
Run
|
Name, grouping, tags, description, job type, and configuration sent to W&B. Project and entity come from this tracker. |
required |
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If it is already initialized |
log_metrics(metrics, step=None, prefix=None) Source
¶
Log metrics to wandb.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
metrics
|
dict[str, float]
|
Dictionary of metric names to values |
required |
step
|
int | None
|
Optional step number |
None
|
prefix
|
str | None
|
Optional prefix for metric names |
None
|
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If the tracker is not initialized |
finish(exit_code=0) Source
¶
Finish the wandb run.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
exit_code
|
int
|
Exit code (0 for success, non-zero for failure) |
0
|
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If the tracker is not initialized |
log_summary(summary) Source
¶
Log summary statistics.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
summary
|
dict[str, Any]
|
Summary dictionary |
required |
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If the tracker is not initialized |
log_model(model, name='model', *, aliases=None, metadata=None) Source
¶
Log a model checkpoint as a W&B artifact. Args: model: Trained PyTorch model name: Name of the model artifact aliases: Optional list of aliases for the artifact metadata: Optional metadata dictionary for the artifact Raises: RuntimeError: If the tracker is not initialized