Input shape verification for convolutional models

[MicheleCattaneo]

What does this PR do?

It introduces utilities to handle invalid input shapes for convolutional models.

Current progress

• Functions to check the input shapes for UNet, HalfUNet and CustomUNet + tests.
• Padding and un-padding functions to be used in conjunction with the aforementioned ones + tests.
• Automatic padding for torch lightning modules exposed through their settings.

Related Issue

closes #9

This draft PR will be updated as the work progresses. Please share any feedback on the overall structure.

[LBerth]

Hi !
Very good job and very nice first contribution.
I have only 2 minor comments concerning the unit test of the padding.
Léa

[colon3ltocard]

Hi again @MicheleCattaneo ,

Sorry for the late reply we went thru a lot of changes before being able to integrate/discuss new features.

We discussed with the team and we have some suggestions to make for the autopad you propose:

• Separation Of Concern: auto-padding is only meaningful for a subset of the models so it should not "contaminate" outside the subset of models (the actual ModelABC subclasses) it is related to.

We propose that a model able to autopad should expose it thru extra settings, here is an example with HalfUnet:

@dataclass_json
@dataclass(slots=True)
class HalfUNetSettings:
    num_filters: int = 64
    dilation: int = 1
    bias: bool = False
    use_ghost: bool = False
    last_activation: str = "Identity"
    absolute_pos_embed: bool = False
    autopad: bool = False

Now the model can handle this in its forward method (this code could be factored in a common class or Mixin):

def forward(self, x):
     if self.settings.autopad:
          x = self.pad(x)
          y=self._forward(x)
          return self.unpad(y)
     else:
        return self._forward(x)

Note : i omitted the shape tracking on purpose for the brevity of this.

This way the autopad doesn't "pollute" the Segmentation/Regression lightning, we don't have to check if autopad is supported using dedicated code outside the model and we can also use it transparently in pure PyTorch loops.

[MicheleCattaneo]

Hi @colon3ltocard,
I agree with what you suggest. Commit 337ef25 implements a first variant of this idea. Let me know if you have any feedback!

[colon3ltocard]

@MicheleCattaneo A few minor comments + some CI fails and then we can merge 👍

[colon3ltocard]

One last comment: directly reference the self._settings for the autopad_enabled bool

[MicheleCattaneo]

Done 👍

[LBerth]

Hi ! Very good contribution, thank you Michele ! :)