W&B Support for the HF Flax Community Week

Flax/JAX Projects
1

Hey all!

I’m super stoked be be taking part in another community week (project “BART pre-training” woop woop!).

I work at Weights & Biases and to help out, we are delighted to support the community here in 3 ways:

  • Free W&B Team account for all projects
  • W&B setup/Q&A sessions with all HF Projects for the duration of the event
  • Code to log to W&B

Full information is here: W&B HF Community Week Suppport

W&B Team sign-up

Add your HF Discord username and Project Name to this Google Sheet and I’ll ping you when its set up!

Logging with W&B

Below is all you need to get rich logging to your W&B workspace. W&B can pick up Tensorboard log files so there is no need to add additional logging calls, unless you’d like to log additional metrics not logged to Tensorboard. See our docs for the W&B Tensorboard integration

First, install wandb

pip install wandb --upgrade

Then use the code below to log to weights and biases

def main():

  import wandb
  if jax.process_index() == 0
    wandb.init(
      entity='my_team_name_or_username', 
      project='norwegian-t5-demo',
      sync_tensorboard=True
  )  

  # log your configs with wandb.config, accepts a dict
  if jax.process_index() == 0
    wandb.config.update(training_args)  # optional, log your configs
    wandb.config.update(model_args)  # optional, log your configs
    wandb.config.update(data_args)   # optional, log your configs

    wandb.config['my_other_thing'] = 12345  # log additional things

  ...

  # Log additional metrics not logged to Tensorboard
  for b in batches
    ...
    metric = get_metric()
    if jax.process_index() == 0
      wandb.log({'my_metric' : metric})

Let me know if there are any W&B questions I can help with, and I hope you all enjoy a fantastic week!

Morgan

4 Likes