Save confusion matrix image to Tensorboard

vision
1

Hi everyone,
let’s suppose I have this simple code that creates a confusion matrix:

import torch
from sklearn.metrics import confusion_matrix

output = torch.randn(1, 2, 4, 4)
pred = torch.argmax(output, 1)
target = torch.empty(1, 4, 4, dtype=torch.long).random_(2)
conf_mat = confusion_matrix(pred.view(-1), target.view(-1))

Is there a way to convert conf_mat to an image and save it to Tensorboard ?
I successfully converted it to an image using mlxtend library, but now I found no way to save it Tensorboard using writer.add_image 

import matplotlib.pyplot as plt
from mlxtend.plotting import plot_confusion_matrix

fig, ax = plot_confusion_matrix(conf_mat=conf_mat)
plt.show()

fig is of type matplotlib.figure.Figure and it should be converted to a torch.tensor or numpy.array since writer.add_image expects one of these as input. I tried some workarounds but none of them worked.
Any suggestions?

2

Would it be possible to pass conf_mat directly to writer.add_image?
If this method expects a tensor or numpy array, I guess it would be possible or did you get an error?

3

You should use writer.add_figure instead

fig = plt.figure()
plt.imshow(conf_mat)

writer.add_figure('my_image', fig)
4

Could you provide me an example? I tried but it does not work(the code runs, but on Tensorboard there are no images)

5

It does not work, I guess the expected tensor should be a tensor representing an image, while conf_mat is a “raw” confusion matrix, e.g.

[[21, 1], 
 [3, 1]]
6 
# pip install --upgrade torch
# pip install tensorboard

import matplotlib.pyplot as plt

import tensorboard
import torch

from sklearn.metrics import confusion_matrix
from torch.utils.tensorboard import SummaryWriter

print(torch.__version__) # 1.5.1
print(tensorboard.__version__)  # 2.2.2

writer = SummaryWriter('runs/experiment')

output = torch.randn(1, 2, 4, 4)
pred = torch.argmax(output, 1)
target = torch.empty(1, 4, 4, dtype=torch.long).random_(2)
conf_mat = confusion_matrix(pred.view(-1), target.view(-1))

fig = plt.figure()
plt.imshow(conf_mat)
plt.show()

writer.add_figure('my_fig', fig)

And here’s the tensorboard snapshot

Screen Shot 2020-06-24 at 11.45.25 AM
Screen Shot 2020-06-24 at 11.45.25 AM1275×508 33.4 KB

2 Likes
7

I tried with your code but it does not work, using same versions.
The code runs correctly and the image is correctly showed on video (via plt.show()) but it does not write the figure to Tensorboard. The strange behavior is that if i try with a random image(taken for example from my dataset), now it writes correctly to Tensorboard. I can’t see where the problem might be.
Anyway, thank you so much for your patience!

8

Maybe, a simple refresh of the page/board will help :wink: ? It was really that simple in my case. I hope you figure out the issue soon. I am sorry I couldn’t be much helpful.

9

In case it might be useful for someone, by modifying the previous code in this way works in my case (the previous code worked on Colab but not locally in my case)

import matplotlib.pyplot as plt
plt.switch_backend('agg')

import tensorboard
import torch

from sklearn.metrics import confusion_matrix
from torch.utils.tensorboard import SummaryWriter

print(torch.__version__) # 1.5.1
print(tensorboard.__version__)  # 2.2.2

writer = SummaryWriter('runs/experiment')

output = torch.randn(1, 2, 4, 4)
pred = torch.argmax(output, 1)
target = torch.empty(1, 4, 4, dtype=torch.long).random_(2)
conf_mat = confusion_matrix(pred.view(-1), target.view(-1))

fig = plt.figure()
plt.imshow(conf_mat)

writer.add_figure('my_fig', fig)
writer.close()
Now the confusion matrix is correctly written on Tensorboard.
3 Likes