diff --git a/monai/inferers/utils.py b/monai/inferers/utils.py index de53108d1d1..3bfeb60297e 100644 --- a/monai/inferers/utils.py +++ b/monai/inferers/utils.py @@ -319,7 +319,12 @@ def sliding_window_inference( else: output_image_list[0][o_slice_idx] += sw_device_buffer[0].to(device=device) else: - sw_device_buffer[ss] *= w_t + # Predictor outputs may alias input data or each other. + if sw_device_buffer[ss].dtype == w_t.dtype: + sw_device_buffer[ss] = sw_device_buffer[ss] * w_t + else: + sw_device_buffer[ss] = sw_device_buffer[ss].clone() + sw_device_buffer[ss] *= w_t sw_device_buffer[ss] = sw_device_buffer[ss].to(device) _compute_coords(unravel_slice, z_scale, output_image_list[ss], sw_device_buffer[ss]) sw_device_buffer = [] diff --git a/tests/inferers/test_sliding_window_inference.py b/tests/inferers/test_sliding_window_inference.py index 70ebb61639f..2b42d97e325 100644 --- a/tests/inferers/test_sliding_window_inference.py +++ b/tests/inferers/test_sliding_window_inference.py @@ -18,7 +18,7 @@ import torch from parameterized import parameterized -from monai.data.utils import list_data_collate +from monai.data.utils import compute_importance_map, list_data_collate from monai.inferers import SlidingWindowInferer, SlidingWindowInfererAdapt, sliding_window_inference from monai.inferers.utils import _compute_coords from monai.utils import optional_import @@ -70,6 +70,45 @@ class TestSlidingWindowInference(unittest.TestCase): + def test_predictor_weighting_preserves_dtype_rounding(self): + for device, input_dtype, predictor_dtype in itertools.product( + _devices[0], (torch.float16, torch.float32, torch.float64), (torch.float16, torch.float32, torch.float64) + ): + with self.subTest(device=device, input_dtype=input_dtype, predictor_dtype=predictor_dtype): + image = torch.linspace(0.1, 0.9, 125, dtype=input_dtype, device=device).reshape(1, 1, 5, 5, 5) + weights = compute_importance_map((5, 5, 5), mode="gaussian", dtype=input_dtype, device=device) + expected = image.to(predictor_dtype).clone() + expected *= weights + expected = expected.to(input_dtype) / weights + result = sliding_window_inference( + image, (5, 5, 5), 1, lambda patch: patch.to(predictor_dtype).clone(), mode="gaussian" + ) + torch.testing.assert_close(result, expected, rtol=0, atol=0) + + def test_predictor_output_alias_preserves_input(self): + for device, mode, batch_size, structure in itertools.product( + _devices[0], ("constant", "gaussian"), (1, 2), ("tensor", "tuple", "dict") + ): + with self.subTest(device=device, mode=mode, batch_size=batch_size, structure=structure): + image = torch.ones((1, 1, 9, 11, 13), device=device) + original = image.clone() + + def predictor(patch): + if structure == "tuple": + return patch, patch + if structure == "dict": + return {"first": patch, "second": patch} + return patch + + with torch.inference_mode(): + result = sliding_window_inference(image, (5, 5, 5), batch_size, predictor, overlap=0.5, mode=mode) + torch.testing.assert_close(image, original, rtol=0, atol=0) + outputs = ( + result.values() if isinstance(result, dict) else result if isinstance(result, tuple) else (result,) + ) + for output in outputs: + torch.testing.assert_close(output, original) + @parameterized.expand(BUFFER_CASES) def test_buffers(self, size_params, buffer_steps, buffer_dim, device_params): def mult_two(patch, *args, **kwargs):