⚖️ 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 via WandbStepCache.

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:
  • project (str) – The W&B project to use for the workspace.

  • entity (Optional[str], default: None) – The W&B entity (user or organization account) to use for the workspace.

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 the metadata of the Step 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 by WandbWorkspace. 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:
  • project (str) – The W&B project to use.

  • entity (str) – The W&B entity (user or organization account) to use.

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 the TorchTrainStep that logs training and validation metrics to W&B.

This can be used with any Workspace implementation, including WandbWorkspace.

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. The group and name parameters will also automatically be set, so a ConfigurationError will be raised if any of project, entity, group, or name 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.

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 the FlaxTrainStep that logs training and validation metrics to W&B.

This can be used with any Workspace implementation, including WandbWorkspace.

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. The group and name parameters will also automatically be set, so a ConfigurationError will be raised if any of project, entity, group, or name 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.