U-NET architecture: Get a feature map a few layers before the final classification layer.

[muneebafzal]

I wanted to ask if we can get hold of a layer (feature vector before the final layer) and discard the last part of the architecture. In other words, can we encode the input image into an encoded vector using a partial U-Net?
For example, in the image below, can I get handle of the following layers:

[Nic-Ma]

Hi @muneebafzal ,

Thanks for your interest here.
I think you can try to write a simple PartialUNet class which inherits from MONAI BasicUNet:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/basic_unet.py#L228

from monai.networks.nets import BasicUNet


class PartialUNet(BasicUNet):
    def forward(self, x: torch.Tensor):
        x0 = self.conv_0(x)

        x1 = self.down_1(x0)
        x2 = self.down_2(x1)
        x3 = self.down_3(x2)
        x4 = self.down_4(x3)

        u4 = self.upcat_4(x4, x3)
        u3 = self.upcat_3(u4, x2)
        u2 = self.upcat_2(u3, x1)
        u1 = self.upcat_1(u2, x0)
        return u1

Thanks.