Voxel-wise importance weighting for training loss (eg, distance to boundary)

[Thibescobar]

Hello, first of all thank you very much for your excellent work!

I'm wondering if there is an easy or not too tricky way to add a weight map to the data during training? To give more importance to some voxels than others, for example borders, small structures, or crucial regions in application cases.

In Ronneberger's original paper for example, high weights are assigned to pixels located in areas where objects touch, thanks to distance-transform-based weighting.

This would also be interesting in tractography, where large vessels are easier to segment than small ones, but also more present, which could induce a convergence favoring large over small vessels.

Thank you. Cheers.

[Karol-G]

Hey,

Something like this is not implemented in nnU-Net. However, what you could do is to provide the weight map as another input channel to the nnU-Net during training. This way the nnU-Net will automatically learn how to best utilize the weights.

Best, Karol

[karllandheer]

Hello, I see in another thread there was some development on this. Did this go anywhere? I don't see how providing a weight map as input will result in any improved performance on the boundaries. Rather a weight map must be fed into the loss function in some way. This issue comes up quite often even in the radiological space, as often we only care about the boundaries between things (for example, joint space width)

[Karol-G]

Hey all,

Is this issue still relevant?

Best regards, Karol

[Thibescobar]

Hello, thank you very much for the update with masked loss.

This post is still relevant. I would appreciate a hint to change the masked loss to a weighted one. Instead of giving a mask for ignoring regions, I would like to give a map of importance to weight the loss more or less for the concerned voxels according to this map. I can modify some part of the code myself if needed and propose a request then if it works.

Thanks.

Best regards, have a good day. Thibault

[Karol-G]

Hey,

It appears something like this is currently not supported and you would need to make quite some changes to nnU-Net to make it work. This probably needs to be computed on-the-fly as part of the data augmentation and be accounted for in the loss computation. Besides that I cannot give you any better suggestions.

Best regards, Karol