Is there a limit to how many classes a Unet with specific channels/strides can predict

[dyollb]

I have been gradually increasing the number of classes in my dataset from 16 to 27 to 36 or 40.
The whole time I have left the architecture the same, I use

UNet(
    spatial_dims=3,
    in_channels=num_channels,
    out_channels=num_classes,
    channels=(16, 32, 64, 128, 256),
    strides=(2,2,2,2),
    dropout=0,
    num_res_units=2,
    norm=Norm.BATCH,
    act="MISH",
)

Since I increased the number of classes to 36 or 40 I noticed that some classes start learning very late or not at all (seemingly).

If I look at the last layer I noticed the conversion from [1, 32, 48, 48, 48] to [1, 40, 96, 96, 96] (for num_classes=40):

So this convolution not only upsamples the patch but also increases the feature channels from 32 to 40. Is this why I start seeing multiple classes learning very late or not at all? Kind of like the 40 classes are predicted from a linear combination of 32 basis vectors. Should I increase the channels to (32, 64, 128, 256, 512) to increase the number of features to 64 in this layer?

[KumoLiu]

Hi @dyollb,

The ability of a U-Net (or any neural network) to distinguish between different classes is not necessarily limited by the number of output channels or strides. Rather, it is often determined by the complexity and variability of the data, as well as the capacity of the network itself (depth, width, etc.).
Increasing the number of classes naturally makes the learning task more challenging, as the model needs to learn to distinguish between a larger number of groups. Your observation that some classes start learning very late or not at all is aligned with the intuition that some classes may be more difficult to learn than others. There could be several reasons for this phenomena. For example, some classes might have fewer training samples, or the features of some classes might be harder to distinguish.
Another important point to check is the ratio of the channels and the strides with number of classes. If the number of classes is too high compared to the number of channels or the strides, the model may struggle to learn effective representations for each class, especially if the differences between classes are subtle.

Thanks.

[dyollb]

Thanks @KumoLiu
I tried a larger U-Net with channels=(32, 64, 128, 256, 512), and see a similar behavior (some classes start learning later) but faster learning overall and possibly a higher mean validation dice at the end of training (still running, but the trend looks good).