Question about network of Deep Atlas

[mikami520]

Deep Atlas
I am currently working on the DeepAtlas, as I read paper, the segmentation network used U-Net with max pool layer. However, when I went through the Monai.network.nets.unet, it does not have a choice to add max pool layer. I am wondering if there is a way to implement max pool layer to the Monai U-Net. I am looking forward to your response, thanks!

[Nic-Ma]

Hi @ericspod ,

Could you please help share some information about this question?

Thanks in advance.

[ericspod]

Our BasicUNet class uses maxpool so that is an option. For the general UNet class it's general experience that maxpooling isn't optimal versus using strided convolutions as we do. However if you wanted to change how the behaviour works you would have to override the UNet._get_down_layer method in a subclass of UNet which would call the inherited method with arguments passed except with strides set to 1, but return a nn.Sequential object containing the maxpool layer and the layer returned by the called inherited method. I think that will give you the behaviour you're expecting, my sketch of this is below which I haven't tested at all:

def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_top: bool) -> nn.Module:
    dlayer=super()._get_down_layer(in_channels, out_channels, 1, is_top)
    return nn.Sequential([nn.MaxPool(...), dlayer])

[mikami520]

@ericspod Thanks for your help!!!

[ebrahimebrahim]

Hey if you're by any chance referring to this, then you may also want to check out the tutorial notebook here: https://github.com/Project-MONAI/tutorials/tree/main/deep_atlas

it's not an exact implementation of the paper (e.g. there is no max pooling as you are looking for) but it gets the idea

[mikami520]

I am wondering if I use the exact sequence of channels that VoxelMorph uses like [16,32,32,32,32,32,32,32,32,16,16], will the deep atlas run out of memory during training? Since the VoxelMorph has much bigger registration network compared to the tutorial and train the both registration and segmentation network alternately: https://github.com/Project-MONAI/tutorials/tree/main/deep_atlas

[ebrahimebrahim]

You can definitely try that sequence of channels. Whether it runs out of memory depends on how much memory you have. It doesn't hurt to try, but I think that sequence of channels will need a lot of memory. You may also have to adjust the variable resize in the notebook to be sufficiently divisible by 2. Choosing a smaller size is one way to save on memory, but results will not be good if the spatial resolution is too low.

[mikami520]

You can definitely try that sequence of channels. Whether it runs out of memory depends on how much memory you have. It doesn't hurt to try, but I think that sequence of channels will need a lot of memory. You may also have to adjust the variable resize in the notebook to be sufficiently divisible by 2. Choosing a smaller size is one way to save on memory, but results will not be good if the spatial resolution is too low.

@ebrahimebrahim I already had a try and it did run out of memory for the RTX 3090 GPU with 128x128x96 scans