Access UNet bottleneck to add a second head for another loss calculation

[hellebe]

Hey folks,

I have a question regarding the UNet network for segmentation purposes. I want to add a second head to the UNet for a second loss calculation. The UNet parameters for the are the following:

net = UNet( dimensions=2, in_channels=3, out_channels=2, channels=[32, 64, 128, 256, 1024], strides=[2, 2, 2, 2], act=monai.networks.layers.Act.LEAKYRELU )

I want to access the bottleneck layer to put a projection head on it. If I get it right the UNet is build from the bottom layer recursively to the top with several nn.Sequential blocks. This makes it hard to iterate over the layers. In a hard coded network I could simply go for something like net.layer[bottleneck] to use the output for further loss calculations.
In the case of the UNet this is not possible. The UNet model consists of three blocks. net.model[0] is the first layer/sequential, net.model[2] is the last layer/sequential.
Going for net.model[1] just gives me a long Sequential again which I cannot iterate over to get the bottleneck layer.
Did someone already have this problem and was able to solve it or find a work around? I hope I just don't miss the obvious solution...

Tl;dr: Access the tensor/layer of the bottleneck in the MONAI UNet for further use.

First discussion btw. I hope I could explain my problem properly.
Thanks in advance :)

[KumoLiu]

Hi @hellebe,
Since both UNet and DynUNet were rewritten in a recursive manner. It may not be easy to get the bottom latent layer.
Here, I showed an ugly way to access the bottom bottleneck. Hope it may help.

from monai.networks.nets import UNet

net = UNet(spatial_dims=3, in_channels=1, out_channels=2, channels=(32,64,128,256), strides=(2,2,2))

#! hard code
print(net.model[1].submodule[1].submodule[1].submodule)

#! iterate it
module_names  = ["model",] + ["submodule"] * 3
for name in module_names[:-1]:
    net = net.get_submodule(name)[1]
print(net.submodule)

[hellebe]

Hey @KumoLiu,

thanks for your help! This way is completely sufficient for my purposes :)
Have a great day.