⚖️ Weights & Biases#
Important
To use this integration you should install tango
with the “wandb” extra
(e.g. pip install tango[wandb]
) or just install the wandb
library after the fact
(e.g. pip install wandb
).
Components for Tango integration with Weights & Biases.
Overview#
The main components provided by this integration are the WandbWorkspace
and
the WandbTrainCallback
.
The WandbWorkspace
is a Workspace
implementation that is
great for collaboration. It tracks Tango runs and steps in the W&B project of your choosing
and uses W&B Artifacts to cache step results in the cloud so that they’re accessible anywhere.
And if you’re training PyTorch models via the TorchTrainStep
,
you can use the WandbTrainCallback
to track metrics throughout the run.
Reference#
- class tango.integrations.wandb.WandbWorkspace(project, entity=None)[source]#
This is a
Workspace
that tracks Tango runs in a W&B project. It also stores step results as W&B Artifacts viaWandbStepCache
.Each Tango run with this workspace will generate multiple runs in your W&B project. There will always be a W&B run corresponding to each Tango run with the same name, which will contain some metadata about the Tango run. Then there will be one W&B run for each cacheable step that runs with a name corresponding to the name of the step. So if your Tango run includes 3 cacheable steps, that will result in a total of 4 new runs in W&B.
- Parameters:
Tip
Registered as a
Workspace
under the name “wandb”.Tip
If you want to change the artifact kind for step result artifacts uploaded to W&B, add a field called
artifact_kind
to themetadata
of theStep
class.This can be useful if you want model objects to be added to the model zoo. In that you would set
artifact_kind = "model"
. For example, your config for the step would look like this:{ type: "trainer", step_metadata: { artifact_kind: "model" }, ... }
Or just add this to the
METADATA
class attribute:@Step.register("trainer") class TrainerStep(Step): METADATA = {"artifact_kind": "model"}
- class tango.integrations.wandb.WandbStepCache(project, entity)[source]#
This is a
StepCache
that’s used byWandbWorkspace
. It stores the results of steps on W&B as Artifacts.It also keeps a limited in-memory cache as well as a local backup on disk, so fetching a step’s resulting subsequent times should be fast.
- Parameters:
Tip
Registered as
StepCache
under the name “wandb”.
- class tango.integrations.wandb.WandbTrainCallback(*args, project=None, entity=None, group=None, name=None, notes=None, tags=None, watch_model=False, wandb_config=None, **kwargs)[source]#
A torch
TrainCallback
for use with theTorchTrainStep
that logs training and validation metrics to W&B.This can be used with any
Workspace
implementation, includingWandbWorkspace
.Tip
Registered as a
TrainCallback
under the name “wandb::log”.Important
When this callback is used with the
WandbWorkspace
it will log metrics to the same W&B project that the workspace uses. Thegroup
andname
parameters will also automatically be set, so aConfigurationError
will be raised if any ofproject
,entity
,group
, orname
are set in this callback.- Parameters:
project (
Optional
[str
], default:None
) – W&B project to associated this run with.entity (
Optional
[str
], default:None
) – W&B entity (user or organization) to associated this run with.group (
Optional
[str
], default:None
) – W&B group to associated this run with.name (
Optional
[str
], default:None
) – Set the name of the run in W&B. If not set, the default will be the name of the step.notes (
Optional
[str
], default:None
) – Arbitrary notes to add in W&B to this run.tags (
Optional
[List
[str
]], default:None
) – Arbitrary tags to add in W&B to this run.watch_model (
bool
, default:False
) – IfTrue
,wandb.watch()
is called to collect gradients and other information about the model throughout training. See docs.wandb.ai/ref/python/watch.wandb_config (
Optional
[Dict
[str
,Any
]], default:None
) – Arbitrary configuration fields to set in W&B for this run. See docs.wandb.ai/guides/track/config.
- class tango.integrations.wandb.WandbFlaxTrainCallback(*args, project=None, entity=None, group=None, name=None, notes=None, tags=None, watch_model=False, wandb_config=None, **kwargs)[source]#
A flax
TrainCallback
for use with theFlaxTrainStep
that logs training and validation metrics to W&B.This can be used with any
Workspace
implementation, includingWandbWorkspace
.Tip
Registered as a
TrainCallback
under the name “wandb::log_flax”.Important
When this callback is used with the
WandbWorkspace
it will log metrics to the same W&B project that the workspace uses. Thegroup
andname
parameters will also automatically be set, so aConfigurationError
will be raised if any ofproject
,entity
,group
, orname
are set in this callback.- Parameters:
project (
Optional
[str
], default:None
) – W&B project to associated this run with.entity (
Optional
[str
], default:None
) – W&B entity (user or organization) to associated this run with.group (
Optional
[str
], default:None
) – W&B group to associated this run with.name (
Optional
[str
], default:None
) – Set the name of the run in W&B. If not set, the default will be the name of the step.notes (
Optional
[str
], default:None
) – Arbitrary notes to add in W&B to this run.tags (
Optional
[List
[str
]], default:None
) – Arbitrary tags to add in W&B to this run.watch_model (
bool
, default:False
) –If
True
,wandb.watch()
is called to collect gradients and other information about the model throughout training. See docs.wandb.ai/ref/python/watch.wandb_config (
Optional
[Dict
[str
,Any
]], default:None
) –Arbitrary configuration fields to set in W&B for this run. See docs.wandb.ai/guides/track/config.