Search⌘ K
AI Features
Home
Courses
Deep Learning with JAX and Flax

How to Use TensorBoard with Flax

Explore how to use TensorBoard with Flax by logging images, text, and training metrics. Understand how to visualize data such as sample images and predictions during model training to better monitor and analyze deep learning workflows.

We'll cover the following...

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.

Python 3.8
from torch.utils.tensorboard import SummaryWriter
import torchvision.transforms.functional as F
logdir = "runs"
writer = SummaryWriter(logdir)
def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
writer.flush()

In the code above:

  • Lines 1–2: We import the SummaryWriter module from torch.utils.tensorboard to log ...