How to use Tensorboard with PyTorch

Let's directly dive in. The thing here is to use Tensorboard to plot your PyTorch trainings. For this, I use TensorboardX which is a nice interface communicating Tensorboard avoiding Tensorflow dependencies.

First install the requirements;

1	pip install tensorboard

2	pip install tensorboardX

Things thereafter very easy as well, but you need to know how you need to communicate with the board to show your training and it is not that easy, if you don't know Tensorboard hitherto.

01

...

02	from tensorboardX import SummaryWriter

03

...

04

05	writer = SummaryWriter('your/path/to/log_files/')

06

07

...

08	# in training loop

09	writer.add_scalar('Train/Loss', loss, num_iteration)

10	writer.add_scalar('Train/Prec@1', top1, num_iteration)

11	writer.add_scalar('Train/Prec@5', top5, num_iteration)

12

13

...

14	# in validation loop

15	writer.add_scalar('Val/Loss', loss, epoch)

16	writer.add_scalar('Val/Prec@1', top1, epoch)

17	writer.add_scalar('Val/Pred@5', top5, epoch)

You can also see the embedding of your dataset

1	from torchvision import datasets

2	from tensorboardX import SummaryWriter

3

4	dataset = datasets.MNIST('mnist', train=False, download=True)

5	images = dataset.test_data[:100].float()

6	label = dataset.test_labels[:100]

7	features = images.view(100, 784)

8	writer.add_embedding(features, metadata=label, label_img=images.unsqueeze(1))

This is also how you can plot your model graph. The important part is to give the output tensor to writer as well with you model. So that, it computes the tensor shapes in between. I also need to say, it is very slow for large models.

01	import torch

02	import torch.nn as nn

03	import torchvision.utils as vutils

04	import numpy as np

05	import torch.nn.functional as F

06	import torchvision.models as models

07	from tensorboardX import SummaryWriter

08

09	class Mnist(nn.Module):

10	def __init__(self):

11	super(Mnist, self).__init__()

12	self.conv1 = nn.Conv2d(1, 10, kernel_size=5)

13	self.conv2 = nn.Conv2d(10, 20, kernel_size=5)

14	self.conv2_drop = nn.Dropout2d()

15	self.fc1 = nn.Linear(320, 50)

16	self.fc2 = nn.Linear(50, 10)

17	self.bn = nn.BatchNorm2d(20)

18	def forward(self, x):

19	x = F.max_pool2d(self.conv1(x), 2)

20	x = F.relu(x)+F.relu(-x)

21	x =F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))

22	x = self.bn(x)

23	x = x.view(-1, 320)

24	x = F.relu(self.fc1(x))

25	x = F.dropout(x, training=self.training)

26	x = self.fc2(x)

27	x = F.log_softmax(x)

28

return x

29

30	model = Mnist()

31

32	# if you want to show the input tensor, set requires_grad=True

33	res =model(torch.autograd.Variable(torch.Tensor(1,1,28,28), requires_grad=True))

34

35	writer = SummaryWriter()

36	writer.add_graph(model, res)

37

38	writer.close()