How to Use TensorBoard with Flax

With TensorBoard installed and some basics out of the way, let’s look at how we can use it in Flax. Let’s use the SummaryWriter from PyTorch to write to the log folder.

How to log images with TensorBoard in Flax

We may want to log sample images when solving computer vision problems. We can also log predictions while training the model. For example, we can log prediction images containing bounding boxes for an object detection network.

Let’s look at how we can log an image to TensorBoard.