fix: use per-sample dice computation in DiceLoss to match CrossEntropyLoss mean reduction (#74)

[rtmalikian]

Fixes #74

Problem

DiceLoss._dice_loss() computes the dice coefficient using torch.sum across all batch elements simultaneously, rather than per-sample. While the dice coefficient is a ratio (bounded 0-1), computing it over the entire batch means that samples with different characteristics — common in medical imaging where anatomy and pathology vary across subjects — can mask each other's contribution to the loss.

When combined with CrossEntropyLoss() (which defaults to reduction='mean'), as done in all training scripts:

supervised_loss = 0.5 * (loss_dice + loss_ce)

the dice loss behavior doesn't align with the expected per-sample mean reduction.

Solution

Changed DiceLoss._dice_loss() to compute the dice coefficient per-sample and return the batch average. This ensures:

1. Each sample contributes equally to the training signal, regardless of other samples in the batch
2. Behavior is consistent with CrossEntropyLoss's default mean reduction
3. No change to the mathematical dice formula — same coefficient, just computed per-sample

Verification

Tested with:

• Basic functionality (2-class, 4-class, 9-class segmentation)
• Gradient flow (no NaN/Inf gradients)
• Various batch sizes (1, 2, 4, 8)
• Per-sample consistency: same sample gives identical loss whether processed alone or in a batch
• Different batch composition: correctly captures per-sample loss when batch samples have different characteristics (the old approach could mask poorly-predicted samples)

=== Test 1: Basic functionality ===
Dice loss: 0.3883
PASS

=== Test 2: Gradient flow ===
Gradient norm: 0.005328
PASS

=== Test 3: Different batch sizes ===
  batch_size=1: loss=0.3896
  batch_size=2: loss=0.3815
  batch_size=4: loss=0.3859
  batch_size=8: loss=0.3893
PASS

=== Test 4: Multi-class ===
  n_classes=2: loss=0.3839
  n_classes=4: loss=0.6346
  n_classes=9: loss=0.8206
PASS

=== Test 5: Per-sample consistency ===
  Single sample loss: 0.382276
  Duplicated batch loss: 0.382276
PASS

=== All tests passed ===

---

About the Author: Raphael Malikian — Clinical AI Solutions Architect. I specialise in building and fixing AI/ML systems for healthcare, including vector databases, RAG pipelines, and clinical NLP. If you need help with your project or think I can add value to your organisation, feel free to reach out — I'd love to connect.

📧 rtmalikian@gmail.com
🔗 GitHub: https://github.com/rtmalikian
🔗 LinkedIn: http://www.linkedin.com/in/raphael-t-malikian-mbbs-bsc-hons-71075436a

---

Disclosure: This code was developed with assistance from mimo-2.5-pro (Xiaomi) via Hermes Agent (Nous Research). All changes were reviewed, tested against the actual codebase, and verified for correctness.