From d0710d9adfd7a96a61d50ca64063986d2fbd6dc6 Mon Sep 17 00:00:00 2001 From: tritsystem Date: Fri, 16 Oct 2026 01:59:38 -0700 Subject: [PATCH] fix(losses): bin each sample from its own range in b-spline GlobalMutualInformationLoss parzen_windowing_b_spline took torch.max/torch.min over the whole batch, so a sample's loss depended on the intensity range of the other images in the batch. Use a per-sample amax/amin, as ITK's Mattes MI does for a single image pair. Batch-of-1 results are unchanged. Assisted-by: Claude Opus 5.5 Signed-off-by: tritsystem --- monai/losses/image_dissimilarity.py | 6 +++++- .../test_global_mutual_information_loss.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/monai/losses/image_dissimilarity.py b/monai/losses/image_dissimilarity.py index 47cab1a2304..e55bf9b82d5 100644 --- a/monai/losses/image_dissimilarity.py +++ b/monai/losses/image_dissimilarity.py @@ -295,7 +295,11 @@ def parzen_windowing_b_spline(self, img: torch.Tensor, order: int) -> tuple[torc compute_img = ( img.float() if not torch.is_floating_point(img) or img.dtype in (torch.float16, torch.bfloat16) else img ) - _max, _min = torch.max(compute_img), torch.min(compute_img) + # Bin each image from its own intensity range (as in ITK's Mattes MI): a batch-wide min/max would make one + # sample's loss depend on the other images in the batch. + reduce_dims = tuple(range(1, compute_img.ndim)) + _max = torch.amax(compute_img, dim=reduce_dims, keepdim=True) + _min = torch.amin(compute_img, dim=reduce_dims, keepdim=True) padding = 2 interior_bins = self.num_bins - 2 * padding value_range = _max - _min diff --git a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py index 2857b5285d0..170a304cfb6 100644 --- a/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py +++ b/tests/losses/image_dissimilarity/test_global_mutual_information_loss.py @@ -271,6 +271,23 @@ def test_b_spline_nonzero_ranges_are_finite(self, case_name, dtype, maximum): self.assertIsNotNone(pred.grad) self.assertTrue(torch.isfinite(pred.grad).all()) + def test_b_spline_batch_samples_are_independent(self): + """Verify each sample is binned from its own intensity range, not the batch's. + + A per-sample loss must not change when other images with a different intensity range share the batch. + """ + generator = torch.Generator().manual_seed(0) + pred_a = torch.rand(1, 1, 32, 32, generator=generator, dtype=torch.float64) + target_a = 0.8 * pred_a + 0.2 * torch.rand(1, 1, 32, 32, generator=generator, dtype=torch.float64) + pred_b = 10 * torch.rand(1, 1, 32, 32, generator=generator, dtype=torch.float64) + target_b = 0.8 * pred_b + 2 * torch.rand(1, 1, 32, 32, generator=generator, dtype=torch.float64) + loss = GlobalMutualInformationLoss(kernel_type="b-spline", reduction="none") + + alone = torch.cat([loss(pred_a, target_a).flatten(), loss(pred_b, target_b).flatten()]) + batched = loss(torch.cat([pred_a, pred_b]), torch.cat([target_a, target_b])).flatten() + + np.testing.assert_allclose(batched.numpy(), alone.numpy(), rtol=1e-6) + class TestGlobalMutualInformationLossBuffers(unittest.TestCase): def test_gaussian_kernel_registers_buffers(self):