Skip to content

Fix class imbalance in CE loss#56

Merged
michaelmckinsey1 merged 11 commits intoLBANN:mainfrom
PatrickRMiles:miles30/class_imbalance
May 6, 2026
Merged

Fix class imbalance in CE loss#56
michaelmckinsey1 merged 11 commits intoLBANN:mainfrom
PatrickRMiles:miles30/class_imbalance

Conversation

@PatrickRMiles
Copy link
Copy Markdown
Collaborator

@PatrickRMiles PatrickRMiles commented Apr 22, 2026 •
edited
Loading

This PR adds class weights to our CE loss calculation to counteract the significant overrepresentation of empty/background voxels in each sample.

Needs:

  • If we get unlucky and choose a subset of samples which are missing one or more classes, we should not end up downweighting the missing class(es). Every non-background class should probably get the same final weight.
  • Testing to ensure this has the intended effect of speeding up convergence

@PatrickRMiles PatrickRMiles changed the title Draft: Miles30/class imbalance Fix class imbalance in CE loss Apr 22, 2026
@PatrickRMiles PatrickRMiles force-pushed the miles30/class_imbalance branch from d035b25 to f2d1f15 Compare April 22, 2026 23:26
@PatrickRMiles PatrickRMiles changed the title Fix class imbalance in CE loss Draft: Fix class imbalance in CE loss Apr 22, 2026
@PatrickRMiles PatrickRMiles marked this pull request as draft April 23, 2026 18:06
@PatrickRMiles
Copy link
Copy Markdown
Collaborator Author

PatrickRMiles commented May 6, 2026

In my testing, it was not clear that this speeds up convergence. However, this fix is still worth implementing.

@PatrickRMiles PatrickRMiles force-pushed the miles30/class_imbalance branch from f2d1f15 to 81dc5bc Compare May 6, 2026 15:15
@PatrickRMiles PatrickRMiles marked this pull request as ready for review May 6, 2026 15:16
@PatrickRMiles PatrickRMiles self-assigned this May 6, 2026
@PatrickRMiles PatrickRMiles force-pushed the miles30/class_imbalance branch from ec57898 to 1ea83f6 Compare May 6, 2026 15:18
@PatrickRMiles PatrickRMiles changed the title Draft: Fix class imbalance in CE loss Fix class imbalance in CE loss May 6, 2026
michaelmckinsey1
michaelmckinsey1 requested changes May 6, 2026
View reviewed changes
Comment thread ScaFFold/utils/config_utils.py Outdated
Comment thread ScaFFold/utils/trainer.py Outdated
Comment thread ScaFFold/configs/benchmark_default.yml Outdated
Comment thread ScaFFold/utils/trainer.py Outdated
PatrickRMiles and others added 2 commits May 6, 2026 10:44
@PatrickRMiles @michaelmckinsey1
remove default ce_weight_num_samples
42ab124 
Co-authored-by: Michael McKinsey <MichaelMcKinsey1@gmail.com>
@PatrickRMiles @michaelmckinsey1
remove default ce_weight_num_samples in trainer
b289d65 
Co-authored-by: Michael McKinsey <MichaelMcKinsey1@gmail.com>
@michaelmckinsey1 michaelmckinsey1 self-requested a review May 6, 2026 18:19
@michaelmckinsey1 michaelmckinsey1 merged commit 6b252a2 into LBANN:main May 6, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@michaelmckinsey1 michaelmckinsey1 michaelmckinsey1 approved these changes

Assignees

@PatrickRMiles PatrickRMiles

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@PatrickRMiles @michaelmckinsey1