How to add model tensors to TensorBoard histogram

TensorBoard supports histograms, which you can use to showcase your tensors. Helpful to see the change in the weights over the training.

To get the tensors from your model you can check the network.state_dict() for the named tensor. Then we can save that tensor into the histogram.

from torch.utils.tensorboard import SummaryWriter

sw = SummaryWriter()

network = CNN()
dataloader = ["sample data"]

for batch in dataloader:
    network.train()

    for param_tensor in network.state_dict():
        sw.add_histogram(param_tensor, network.state_dict()[param_tensor], global_step=global_step)

This will save a bucketed version of the tensors. See arguments for add_histogram for more info.

Note: This will save a large amount of data. Consider which tensors you may want to add as a histogram.