Feature: Add independent channel option in UNet – Conv group implementation

[melisande-c]

Feature

Option to conveniently train independent UNets for each channel in parallel. This is useful for noise2void, most cases will have channel independent noise.

Implementation

PyTorch has the option to pass the parameter groups to Conv2d & Conv3d. groups=n creates n parallel convolutions with no shared connections; in this case we have groups==in_channels and out_channels==K*in_channels – which is also known as a "depthwise convolution". Using depthwise convolutions throughout the network creates essentially n=in_channels parallel networks.

When concatenating the skip connections in the decoder we have to interleave the input and the encoder features to ensure no connections between the parallel networks.

No change has to be made to batch normalisation because it is channel-wise.

Change details

• add independent_channels option to UNet class and UNetModel
• add groups parameter to UnetEncoder and UnetDecoder; which will be passed to each of their ConvBlock layers.
• if independent_channels=True set groups==in_channels in the encoder and decoder. set groups=1 otherwise.
• add _interleave static method in decoder for interleaving input and skip connections.
• update tests to cover independent channels at different depths
• add test to validate channels are indeed independent throughout network
• !possible breaking change! independent_channels default is set to True

[jdeschamps]

Looks good on my side!