Trying to solve issues with UNet and ASPPHeads.

[Joao-L-S-Almeida]

Alternative to #456.

[singhshraddha]

If this is supposed to represent the out_channels per dec_outs item then it's not accurate, it should be:

self.out_channels = [out_channels * 2**i for i in reversed(range(num_stages))]

For out_channels = 32, num_stages = 5, self.out_channels = [512, 256, 128, 64, 32] which corresponds to the channel of items in dec_outs

[Joao-L-S-Almeida]

Thanks. Just filling space to see things running.

[singhshraddha]

I looked into this a bit and setting the whole model to eval also affects the dropout layers, so I don't think we can do this

A better way to just set batch norm layers to eval is mentioned in this post. But, more importantly I don't think we should print a warning about using BatchNorm when batch size is 1 because it is also just a bad idea to do so

[Joao-L-S-Almeida]

I believe that doesn't make much sense to use batch_size=1 for models with batch normalization layers, even if we add a kind of workaround to make it run.
Maybe the best solution is to include a global check, for any model in terratorch, which will raise and Exception for this cases, indicating the user to adopt batch_size > 1 when BatchNorm is used.
What do you think, @romeokienzler ?

[singhshraddha]

This is actually for line 276 but I can only review changed lines

The way that ASPPRegressionHead is setup doesn't comply with the way encoder_decoder_factory.py calls decoder build. Since, the build expects that the first arguments to be channel_list and we have dilations instead. So, the dilations are set incorrectly and I cannot set dilations using a config

To resolve this I'd suggest

• not accepting in_channels as input parameter
• setting channel_list as the first input parameter
• Using in_index with channel_list to set the in_channels

class ASPPRegressionHead(ASPPHead):

    def __init__(
        self, 
        channel_list,
        dilations:list | tuple =(1, 6, 12, 18), 
        channels:int=None,
        out_channels:int=1,
        align_corners=False,
        head_dropout_ratio:float=0.3,
        input_transform: str = None,
        in_index: int = -1,
        **kwargs
):

   self.in_channels = channel_list[in_index]

Note this only works when input_transform is None. I currently only care about this use case wherein we select one item from the channel list and this corresponds to the item we pick from the dec_outs. Though if you want to make it more generic mmseg has a BaseDecodeHead._init_inputs function that deals with different types of channel_list, input_transforms etc

[Joao-L-S-Almeida]

Thanks. I'll try to adapt and include it in the ASPP module.

[singhshraddha]

Same as the review in regression_tasks.py for setting model.eval()

[Joao-L-S-Almeida]

I have removed all these calls. I believe the best approach it to block the usage of batch_size < 2 when BatchNorm is involved. I'll do it in another PR.

[singhshraddha]

Looks fine for my use case. I'll run a model next week using this branch thanks