[Feature request] When using PyTorch backend, use torch.data.tensorboard in keras.callbacks.TensorBoard

[foxik]

Hi all,

even with PyTorch backend, keras.callbacks.TensorBoard uses tensorflow.summary to generate the Tensorboard logs, which can be inconvenient when only PyTorch is installed.

However, PyTorch has a module torch.data.tensorboard capable of generating Tensorboard logs. As a very simple PoC, I put together the following simple callback:

class TorchTensorBoard(keras.callbacks.Callback):
    def __init__(self, path):
        self._path = path
        self._writers = {}

    def writer(self, writer):
        if writer not in self._writers:
            import torch.utils.tensorboard
            self._writers[writer] = torch.utils.tensorboard.SummaryWriter(os.path.join(self._path, writer))
        return self._writers[writer]

    def add_logs(self, writer, logs, step):
        for key, value in logs.items():
            self.writer(writer).add_scalar(key, value, step)

    def on_epoch_end(self, epoch, logs=None):
        if logs:
            self.add_logs("train", {k: v for k, v in logs.items() if not k.startswith("val_")}, epoch + 1)
            if isinstance(getattr(self.model, "optimizer", None), keras.optimizers.Optimizer):
                self.add_logs("train", {"learning_rate": self.model.optimizer.learning_rate}, epoch + 1)
            self.add_logs("val", {k[4:]: v for k, v in logs.items() if k.startswith("val_")}, epoch + 1)

which works fine with Keras 3.

I assume the most convenient way would be to support this torch.utils.tensorboard "backend" in keras.callbacks.TensorBoard directly (with some of the features disabled, like profiling and probably write_graph, ...) when torch Keras backend is active.

I would find such an update to keras.callbacks.TensorBoard useful, but I am not volunteering to implement it :thinking:

[qlzh727]

I think under the hood, pytorch tensorboard API also forward the numpy value to the tf.summary API. I think the current.

Does the current Tensorboard callback raise any error at the moment? (or it just need a extra dependency for TF)

[foxik]

I think under the hood, pytorch tensorboard API also forward the numpy value to the tf.summary API. I think the current.

No, it does not -- PyTorch Tensorboard API works without the tensorflow package installed (only tensorboard package needs to be installed).

Does the current Tensorboard callback raise any error at the moment? (or it just need a extra dependency for TF)

With PyTorch Keras backend, when tensorflow dependency installed, the keras.callbacks.TensorBoard does work.

So the whole issue is about allowing keras.callbacks.TensorBoard to work without the tensorflow dependency when PyTorch backend is used.

[nicolas-combaret]

Hello, is there a plan to implement this feature soon?