`log_graph = True` for tensorboard logger doesn't show model graph.

Describe the bug I wanted to show the computational graph of the patchcore model and for that I enable project.logger: tensorboard and set log_graph = True

logger = AnomalibTensorBoardLogger(
            name="Tensorboard Logs",
            save_dir=os.path.join(config.project.path, "logs"),
            log_graph=True
        )

Now, when I finished training, I run the tensorboard and under the GRAPHS tag, it showed nothing.

To Reproduce Steps to reproduce the behavior:

1. Go to the patchcore config file and set `logger: tensorboard’
2. Next, working/anomaly_detection_engine/anomalib/anomalib/utils/loggers/__init__.py and set log_graph=True
3. Run the anomalib with patchcore config
4. After training, run the event log in tensorboard.

Expected behavior

• It should give a computational graph, like https://github.com/PyTorchLightning/pytorch-lightning/issues/2915

Screenshots

• (Mention above)

Hardware and Software Configuration

• OS: [Ubuntu, OD] Ubuntu
• PyTorch-Lighting: ‘1.5.9’
• NVIDIA Driver Version [470.57.02]
• CUDA Version [e.g. 11.4] 10.2
• CUDNN Version [e.g. v11.4.120] 7605
• OpenVINO Version [Optional e.g. v2021.4.2]

Additional context

• I also found that the add_image function of the tensorboard callback, doesn’t write any image on the tensorboard either.
• Is there any way, I can write the Histogram in tensorboard in pytorch-lighting. Histograms are made for weights and bias matrices in the network. They tell us about the distribution of weights and biases among themselves. Is there any convenient way to add it to anomalib’ models?

[ashwinvaidya17]

Edit: I have moved the for loop after trainer.fit so that it works with Patchcore @innat I am working on updating the documentation but for now you can use the following snippet in train.py to log model graph.

for logger in loggers:
    if isinstance(logger, AnomalibWandbLogger):
        # NOTE: log graph gets populated only after one backward pass. This won't work for models which do not
        # require training such as Padim
        logger.watch(model, log_graph=True, log="all")
    elif isinstance(logger, AnomalibTensorBoardLogger):
        logger._log_graph = True
        logger.log_graph(model, input_array=torch.ones((1, 3, 256,256)))

So your train method should look like this

def train():
    """Train an anomaly classification or segmentation model based on a provided configuration file."""
    args = get_args()
    config = get_configurable_parameters(model_name=args.model, config_path=args.config)

    if config.project.seed != 0:
        seed_everything(config.project.seed)

    datamodule = get_datamodule(config)
    model = get_model(config)
    loggers = get_logger(config)

    callbacks = get_callbacks(config)

    trainer = Trainer(**config.trainer, logger=loggers, callbacks=callbacks)
    trainer.fit(model=model, datamodule=datamodule)

    for logger in loggers:
        if isinstance(logger, AnomalibWandbLogger):
            # NOTE: log graph gets populated only after one backward pass. This won't work for models which do not
            # require training such as Padim
            logger.watch(model, log_graph=True, log="all")
        elif isinstance(logger, AnomalibTensorBoardLogger):
            logger._log_graph = True
            logger.log_graph(model, input_array=torch.ones((1, 3, 256,256)))

    # load best model from checkpoint before evaluating
    load_model_callback = LoadModelCallback(weights_path=trainer.checkpoint_callback.best_model_path)
    trainer.callbacks.insert(0, load_model_callback)

    trainer.test(model=model, datamodule=datamodule)

And here is the output

Not sure if this is what you wanted. Also, to access other methods provided by tensorboard you can access logger.experiment object. But I’d be more inclined towards using this https://pytorch-lightning.readthedocs.io/en/stable/api/pytorch_lightning.loggers.tensorboard.html rather than accessing experiment object. I’ll try to make this more clear in the PR.

[innat]

@ashwinvaidya17 Thanks a lot, it’s really helpful. I’m currently out of my workplace, will check soon.