Loss Function and Non-linearity when using 'DC_and_BCE_loss'

[davidguo123456]

Hi,

I'm adapting nnUNetv2 to add a classification head (cnn-based) for multitask and I was running into issues with my architecture and actually being able to train the model. More specifically, I was able to get classifier losses to decrease only when I removed every single ReLU unit between my conv and fc layers, however this led to poor validation loss results. I then checked the 'DC_and_BCE_loss' implementation and noticed it mentioned a blanket "DO NOT APPLY NONLINEARITY IN YOUR NETWORK!" warning. However, in the docs it mentions that you should avoid linearity at the end of the architecture adjustments. Removing linearity only at the end however still led to a stagnant classifier training and validation loss. My question is, how exactly is nnUNetv2 adding in non-linearity and why does it interfere with my classifier?

Thanks!