Weights and Biases Callback for Elegy - Githubissues

poets-ai / elegy

A High Level API for Deep Learning in JAX

https://poets-ai.github.io/elegy/

MIT License

469 stars 32 forks source link

Weights and Biases Callback for Elegy #220

Closed soumik12345 closed 2 years ago

soumik12345 commented 2 years ago

Callback that streams epoch results to a Weights & Biases run.

Sample Usage:

wandb.login()
wandb_logger = WandbCallback(project="sample-wandb-project", job_type="train")
model.fit(X_train, Y_train, callbacks=[wandb_logger])

Key Features:

Automatically log your experiment results using elegy.callbacks.WandbCallback
Save and version model checkpoints as Weights & Biases Artifacts

Colab Notebook for Demonstration: https://colab.research.google.com/drive/1M-KJcxFx5LXVqtt7Yztbf8uyX0Cd7GEw?usp=sharing Sample Dashboard: https://wandb.ai/geekyrakshit/elegy-test

cgarciae commented 2 years ago

Hey @soumik12345! Sorry it took so long to review. I updated the dependencies and solved merge issue.

Thanks a lot for doing this ❤️