Visualizing the UNet in Tensorboard

[lenlobo]

Hi, I have a UNet like

class Net(pytorch_lightning.LightningModule):
    def __init__(self):
        super().__init__()
        self._model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=2,
            channels=(16, 32, 64, 128, 256),
            strides=(2, 2, 2, 2),
            num_res_units=2,
            norm=Norm.BATCH,
        )
        self.loss_function = DiceLoss(include_background=True,to_onehot_y=True, softmax=True)
        self.post_pred = AsDiscrete(argmax=True, to_onehot=2)
        self.post_label = AsDiscrete(to_onehot=2)
        self.dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
        self.best_val_dice = 0
        self.best_val_epoch = 0
        self.validation_step_outputs = []

    def forward(self, x):
        return self._model(x)

Is there a way to visualize this? In either Tensorboard using Graphs or something to just print it in the notebook?
I tried the Tensorboard,SummaryWriter way,
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("torchlogs/")
model = Net()
writer.add_graph(model, X)
writer.close()

This throws an error like 'RuntimeError: Net is not attached to a Trainer.'

Thanks!

[KumoLiu]

Hi @lenlobo, you can refer to this link.

Hope it helps, thanks!

[lenlobo]

Hi @KumoLiu ,
I was looking at that link itself, I am using the Unet from Monai, and when I try to visualise it, a runtime error comes 'RuntimeError: Net is not attached to a Trainer.'
I can visualise let's say a basic NN as mentioned in the link but this UNet doesn't work.

[KumoLiu]

Hi @lenlobo,

here are two solutions might help you:

# using lightning's tb logger
writer = pytorch_lightning.loggers.TensorBoardLogger(save_dir="./")
model = Net()
writer.log_graph(model, torch.rand(5, 1, 32, 32, 32))

# use native summary writer but with raw torch module
writer = SummaryWriter("torchlogs/")
writer.add_graph(model._model, torch.rand(5, 1, 32, 32, 32))
writer.close()

Hope it helps!
Regards

[lenlobo]

Hi @KumoLiu ,
The second method worked :) thank you!