switching to tensorboard logger for organelle phenotyping

[mattersoflight]

Hi @alishbaimran @ziw-liu I looked carefully into model logging and profiling yesterday and figured that W&B is not the optimal backend for our projects. Please start using tensorboard logger in your training script on infection_state branch.

If this turns out to be tricky, I will start a PR after Tuesday to help with the transition.

Here are the features we need for the project, and options.

Model visualization: There doesn’t seem to be a way to extract images from W&B logs to illustrate training runs. See this movie for the visualizations we would need to interpret models. We have code to make movies like this tensorboard and W&B doesn’t yet provide an API to do this.

Profiling resources: W&B’s visualization of resources is useful, but resource usage can be profiled with lightning's SimpleProfiler and PytorchProfiler. The PytorchProfiler is more feature complete than W&B.

Profiling gradients: W&B makes it particularly easy to visualize vanishing gradients in a layer, but this can also be done by logging .grad property of tensors to tensorboard. We should look for a convenient API that plays well with lightning. @ziw-liu have you come across any good solutions?

Sharing logs: W&B makes it easy to share logs via URL, but it is only marginally harder to use tensorboard to share logs by pointing to a directory.

Lastly, monetary and time cost: For the active team of 4, W&B pricing structure will cost us the cost of attendance at a conference for one team member over 6-9 months. In addition to $50/user/month, we’d pay $3/100GB/month. If we use W&B for this project, we will fragment virtual staining and organelle phenotyping logs in different systems. If we use W&B for all of our projects, it is expensive.