How to extract intermediate feature maps from monai's densenet ?

[cp2611]

Hello
I need to store the intermediate output of each dense block of monai densenet.
How can I do that?

[ericspod]

One possible way is with forward hooks:

import monai
import torch

net=monai.networks.nets.Densenet121(spatial_dims=2,in_channels=1,out_channels=1)
output_dict={}


def forward_hook(mod, inputs, outputs):
    output_dict[mod]=outputs

    
def add_hooks(mod):
    if isinstance(mod, monai.networks.nets.densenet._DenseLayer):
        mod.register_forward_hook(forward_hook)
    
    
net.apply(add_hooks)
        
result=net(torch.rand(1,1,128,128))

print(output_dict)

After the forward pass through the network output_dict will map every dense layer to its output tensor. You don't necessarily want to be accumulating these so your hook should instead do whatever it is you want the output for, but this code demonstrates how to setup a hook to get access the input and output tensors.