feat: replace bce by focal loss in linknet loss

[charlesmindee]

Following a suggestion by @fg-mindee and @SiddhantBahuguna, this PR replaces the BCE loss by the Focal loss in the linknet loss to increase the recall (imbalanced classes)

Any feedback is welcome!

[codecov]

Codecov Report

Merging #824 (972a2b6) into main (fae4923) will decrease coverage by 0.02%.
The diff coverage is 92.00%.

@@            Coverage Diff             @@
##             main     #824      +/-   ##
==========================================
- Coverage   96.01%   95.98%   -0.03%     
==========================================
  Files         131      131              
  Lines        5019     5033      +14     
==========================================
+ Hits         4819     4831      +12     
- Misses        200      202       +2     

Flag	Coverage Δ
unittests	95.98% <92.00%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
doctr/models/detection/linknet/tensorflow.py	97.67% <91.66%> (-1.06%)	⬇️
doctr/models/detection/linknet/pytorch.py	97.97% <92.30%> (-0.94%)	⬇️
doctr/transforms/modules/base.py	94.59% <0.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fae4923...972a2b6. Read the comment docs.

[fg-mindee]

Thanks, I added a comment on the implementation! Have you tried to check if it improves training perf?

[fg-mindee]

I think we need to address the masking + reduction problem of the dice loss first: once masked, this reduces the tensor to something in 1D. So if one class has 10 times more masked region than another, this will be a problem

I'd suggest doing the following: changing

seg_target[mask].mean()

to

mask = mask.to(dtype=torch.foat32)
# Average on N, H, W
class_loss = (seg_target * mask).sum((0, 2, 3)) / mask.sum((0, 2, 3))
loss = class_loss.mean()

or average it on H, W only before the final mean

[charlesmindee]

I am not sure I understand well here: do you want to remove each ...[mask] occurrence ?

[fg-mindee]

So the difference is the following:

• my_tensor[mask] is 1D tensor with a number of elements = number of True in the mask
• my_tensor * mask.to(dtype=torch.float32) has the shape shape as my_tensor, it only puts zero on elements that are masked out

Now if you perform a reduction operation like mean:

• in the first case, you divide by the number of elements in my_tensor[mask] = mask.sum()
• in the second one, you dive by the number of elements in my_tensor

And this extends to dimension-specific operations, so if we mask it, we lose the separation of the dimensions to get a contiguous tensor in the end. To properly scale the loss, in the first case, this widely increases the contribution of classes with the highest amount of positive mask (makes no difference if there is only a single class). And since we specifically want to help balance the less-frequent classes here, I suggest leveraging that second option with my suggestion above 👍

[fg-mindee]

(and we need to do this for both the dice loss and the focal loss)

it will only make a difference in cases of multi-class and where the mask isn't only True but that might be safer!

Either way, I think we should run a training with the configuration to make sure this yields a positive change 👍

[fg-mindee]

A few corrections in the loss computation and we'll be good to go!

[fg-mindee]

Final cosmetic adjustments and we're good to go 👍

[fg-mindee]

My bad, a few fixes to do on my previous suggestion

[fg-mindee]

Thanks!