How to design CNN with MSELoss?


CNN is mostly taught for image classification, so we use Softmax + Cross Entropy loss to get probabilities. If we need to design a CNN where we want real-value outputs like in Object boundary detection or colorization where we need to compare original color with deduced colors, then how should we design the CNN?

Should output dimensions be same size as original image size for MSE to work? If so, what would be the point of hierarchical representation in CNN. If we make output dimensions of CNN smaller, then we won’t be able to use MSE. Please let me know if you have any suggestions for this.

Currently I am working on colorization problem and MSE loss is in range 200-300 while CE loss starts with 5~6, I am using ConvTranspose2d to scale the image back to original size in the last layer when using MSE.

Thanks in advance,