Log one Image after every layer

[V-Krause]

I am currently working on a project classifying MRI images for Alzheimer's disease. For illustration I would like to log a sample image after each layer so you can see what the different filters/layers have done to the image. I know that you can log images in Tensorboard with logger.experiment, but is there also a way to insert this between each layer? So I would like to have a series of images that starts with the initial image, then outputs one image after each layer and last but not least, outputs the result image. What would I need to do to achieve this?

Here is an excerpt from my work.

`
class AlzheimerClassifier(LightningModule):

def __init__(self, learning_rate):

    super().__init__()

    self.conv1 = nn.Conv3d(1, 16, kernel_size=(3, 3, 3), padding=0)
    self.conv2 = nn.Conv3d(16, 32, kernel_size=(3, 3, 3), padding=0)
    self.conv3 = nn.Conv3d(32, 64, kernel_size=(3, 3, 3), padding=0)
    self.pool = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=2)
    self.fc1 = nn.Linear(64*14*14*14, 3)

    # init loss function
    self.loss = torch.nn.CrossEntropyLoss()
    
    # init learning rate
    self.learning_rate = learning_rate
    
    # init accuracy funktion
    self.acc = torchmetrics.Accuracy()

    # init confusion matrix function
    self.confmat = torchmetrics.ConfusionMatrix(num_classes=3)
    
def forward(self, x):

    x = self.pool(F.relu(self.conv1(x)))
    x = self.pool(F.relu(self.conv2(x)))
    x = self.pool(F.relu(self.conv3(x)))
    x = x.view(-1, 64*14*14*14)
    x = self.fc1(x)
    
    return x

def training_step(self, batch, batch_idx):

    image, label = batch
    preds = model(image)
    
    loss = self.loss(preds, label)
    self.log('training/step_loss', loss)
    
    acc = self.acc(preds, label)
    self.log('training/step_acc', acc)

    return {'loss': loss, 'acc': acc}

def training_epoch_end(self, outputs):
            
    loss = torch.stack([x['loss'] for x in outputs]).mean()
    self.log('training/epoch_loss', loss)
    
    acc = torch.stack([x['acc'] for x in outputs]).mean()
    self.log('training/epoch_acc', acc)
    
    self.logger.experiment.add_scalar("Loss/Train",
                                      loss,
                                      self.current_epoch)

    self.logger.experiment.add_scalar("Accuracy/Train",
                                      acc,
                                      self.current_epoch)
    
    #self.trainer.progress_bar_callback.main_progress_bar.write(
    print(f"Epoch {self.trainer.current_epoch} train_loss={loss.item()} train_acc={acc.item()}")

def configure_optimizers(self):

    optimizer = torch.optim.Adam(self.parameters(), self.learning_rate)
    return optimizer

`

[rohitgr7]

you can register forward hooks:

class LitModel(LightningModule):
    
    def __init__(self, ...):
        super().__init__()
        self.hooks = {}
        ....

    def log_intermediate(self, name):
        def wrapped(module, inp, output):
            self.logger.experiment.whatever_method_logs_image(f'layer_act/{module}, out.detach(), step=self.global_step)
        return wrapped

    def on_train_epoch_start(self):
        self.hook['conv1'] = self.conv1.register_forward_hook(self.log_intermediate('conv1'))
        # same for conv2, ....

    def on_validation_start(self):
        if self.hooks:
            for hook in list(self.hooks):
                self.hooks.pop(hook).remove()

I haven't actually tried this code, but if it doesn't work for you, we can explore why or look for more options.

[V-Krause]

Yeah, right now Im logging 3D data.
My idea is to always take the same 2D slice from the images, so for example always take the middle slice, dependend on shape (e.g. for 128x128x128 [64,:,:]), from the images.
The question is, at which point do I insert that I only take that particular slice?

[rohitgr7]

your output seems 3d, so you can slice that. out[slice_it].detach()

[V-Krause]

Okay, I think the slicing worked.
I would then only have to enter a different number for each image, since the middle slice of the first image is 64 and the middle slice of the last image is 7.

def log_intermediate(self, name):
        def wrapped(module, inp, out):
            self.logger.experiment.add_image(f'layer_act/{module}', out[:,:,10].detach())
        return wrapped

However, this did not fix the error.
AssertionError: size of input tensor and input format are different. tensor shape: (8, 16, 126, 126), input_format: CHW

The problem still seems to be that the image is not taken from the batch, so that it now still has the form [batch, channel, height, width]. But for the function you need, as far as I understood, [Channel, Height, Width], so the CHW which is required.

But shouldn't the .detach() accomplish this?

[rohitgr7]

the error tells the issue. You are logging [batch_size, C, H, W] tensor which contains multiple images, but add_image supports only logging a single image. If you want to log a batch of images, you need to use https://pytorch.org/docs/stable/tensorboard.html#torch.utils.tensorboard.writer.SummaryWriter.add_images

checkout the dataformats and others parameters in API.

[V-Krause]

Okay, thanks, I'll give it a try.
Maybe one more general question. For my work I need only these images from one image and not from all or from a batch. Is there an easy way to do this? So I guess you could also just take a source image and apply the layers to it step by step and then output the images, right? Or does something happen/change with the image when it goes through the model, as opposed to when I do it manually? I know this may not fit the question anymore, but maybe you can answer this question as well.