Adding mesh to Tensorboard not working via TensorboardWriter

[matthiasvestner]

I cannot manage to add meshes to Tensorboard. I have added add_mesh to self.tb_writer_ftns in logger/visualization.py and was expecting the TensorboardWriter just to wrap SummaryWriters' add_mesh function. However every variant I try to call add_mesh from a TensorboardLogger instance is failing.

import unittest
import torch
import logging
import numpy as np

from logger.visualization import TensorboardWriter
from torch.utils.tensorboard import SummaryWriter

def simple_mesh_generator(alpha=0.2, beta=0.1, n=4):
    t = np.arange(0, 2 * np.pi, 2 * np.pi / n)
    polygon_verts = np.array([beta * np.cos(t), beta * np.sin(t), alpha * np.ones(len(t))]).transpose()
    extr_verts = np.array([[0, 0, 0], [0, 0, 1]])
    V = np.concatenate([extr_verts, polygon_verts])
    lower = np.array([[1, 2 + np.mod(k + 1, n), 2 + np.mod(k, n)] for k in range(n)])
    upper = np.array([[0, 2 + np.mod(k, n), 2 + np.mod(k + 1, n)] for k in range(n)])
    T = np.concatenate([lower, upper])

    return V, T

class TestVisualizationTB(unittest.TestCase):
    def setUp(self) -> None:
        self.logdir = 'test_tensorboard'
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)

        self.logger = logging.getLogger('testlogger')
        self.logger.setLevel(logging.INFO)

        self.V, self.T = simple_mesh_generator()
        self.Vs = torch.stack([torch.from_numpy(self.V), torch.from_numpy(self.V)])
        self.Ts = torch.stack([torch.from_numpy(self.T), torch.from_numpy(self.T)])
        self.colors = 100*torch.ones_like(self.Vs)

    def test_visualize_mesh_native(self):
        writer = SummaryWriter(self.logdir)
        writer.add_mesh('Native mesh', self.Vs, faces=self.Ts, colors=self.colors)

    def test_visualize_mesh_wrapper_v1(self):
        writer = TensorboardWriter(self.logdir, self.logger, True)
        writer.add_mesh('Wrapped mesh 1', self.Vs, faces=self.Ts, colors=self.colors)

    def test_visualize_mesh_wrapper_v2(self):
        writer = TensorboardWriter(self.logdir, self.logger, True)
        writer.add_mesh('Wrapped mesh 1', self.Vs, faces=self.Ts)

    def test_visualize_mesh_wrapper_v3(self):
        writer = TensorboardWriter(self.logdir, self.logger, True)
        writer.add_mesh('Wrapped mesh 2', self.Vs)

[matthiasvestner]

test_visualize_mesh_native is working, all others fail

[matthiasvestner]

Solved. add_mesh expects different order of parameters.

if name == 'add_mesh':
   add_data(tag, data, *args, **{**kwargs, 'global_step': self.step})
else:
   add_data(tag, data, self.step, *args, **kwargs)