TensorBoard is a Python language library that can be used to display graphs and visualizations for PyTorch or TensorFlow neural models. I’m not a fan of TensorBoard but some of my colleagues use it often.

I hadn’t looked at TensorBoard in several months, and because things in PyTorch and TensorFlow are moving at hyper speed, I figured I’d take another look at using TensorBoard to visualize a PyTorch trained model.

Adding code to save a TensorBoard visualization is simple.

I don’t get as much information from a model’s diagram as I do from from the model’s code.

I ran into a problem right away. When you install PyTorch, you automatically get TensorBoard as a module inside the torch.utils module. But my demo program couldn’t find TensorBoard until I manually installed it using the command “pip install tensorboard”.

I took an existing PyTorch program I had, where I used a 6-(10-10)-3 neural network to predict an employee’s job type (mgmt, supp, tech) from their sex, age, city (anaheim, boulder, concord), and income. I added code to save a TensorBoard model. The key statements are:

. . .
net = Net().to(device)
. . .
(train network)
. . .
inpt = np.array([[-1, 0.30,  0,0,1,  0.5000]],
  dtype=np.float32)
inpt = T.tensor(inpt, dtype=T.float32).to(device) 
. . .
from torch.utils.tensorboard import SummaryWriter
print("Saving graph representation using tensorboard")
writer = SummaryWriter()
writer.add_graph(net, inpt)
writer.close()

The SummaryWriter() class is the main object. The add_graph() method accepts a PyTorch model and a tensor input. The close() method automatically saves to a .\runs directory by default.

Next, I launched a second shell, and issued the command “tensorboard –logdir=runs” which fired up a server on my machine. Then I launched a browser, and navigated to http://localhost:6006 and voila! There was the visualization of my model.

Alas, to me, the visualization is relatively awkward to use, and really doesn’t give me much information compared to the network definition code:

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()
    self.hid1 = T.nn.Linear(6, 10)  # 6-(10-10)-3
    self.hid2 = T.nn.Linear(10, 10)
    self.oupt = T.nn.Linear(10, 3)

  def forward(self, x):
    z = T.tanh(self.hid1(x))
    z = T.tanh(self.hid2(z))
    z = self.oupt(z)  
    return z

For me, the code is almost always more informative than a diagram. I can see the architecture clearly from the code. There are some exceptions to my preference for code, but a TensorBoard visualization is not one of them.

So, TensorBoard may be a useful tool for some people, but I’m not a fan.

I am a fan of unusual fans.