add wandb support

[tcapelle]

This PR adds wandb logging for the training script 🚀
With this tool, we can compare runs on the dashboard and then get super nice results.
It also automatically logs the apple system metrics like GPU power, CPU power, etc..
I have been also benchmarking the processor recently!

And check the model run here

[tcapelle]

I also wrote a report on my findings:
https://wandb.ai/tcapelle/apple_m1_pro/reports/Deep-Learning-on-the-M1-Pro-with-Apple-Silicon---VmlldzoxMjQ0NjY3

• I would absolutely recommend you re-run your benchmarks on NGC containers, as the performance for Tensorflow can be substantially higher.

[tlkh]

I did use the NGC TF container.

What is the image size you used? From your script, it looks like it defaults to 128x128? That could be why it looks like it runs faster.

[tcapelle]

Nice!
Yeah, I ran on 128x128.