Skip to content

Commit 586dea1

Browse files
chhayankjainpre-commit-ci[bot]ericspod
authored
fix: replace set_track_meta with F.interpolate in RandSimulateLowResolution (#8409) (#8837)
Fixes #8409 ### Description `RandSimulateLowResolution` internally performs a downsample → upsample cycle using two `Resize` transforms. To prevent these from being recorded in the invertible transform stack, the previous implementation temporarily toggled the global `set_track_meta(False)` flag and restored it afterward. This is not thread-safe: in multi-threaded data loading (e.g. `ThreadDataLoader`), another thread calling `get_track_meta()` between the toggle and the restore would silently receive the wrong value, causing incorrect metadata tracking behaviour. Fix: replace the `Resize` transforms with direct `torch.nn.functional.interpolate` calls on a plain tensor obtained via `convert_to_tensor(img, track_meta=False)`. This avoids any global state mutation entirely. Output dtype (float32) and metadata-copy behaviour are preserved from the original implementation. `set_track_meta` is also removed from the import since it is no longer used anywhere in the file. ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [x] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [ ] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. --------- Signed-off-by: chhayankjain <chhayank44@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent eb5e8e1 commit 586dea1

2 files changed

Lines changed: 74 additions & 27 deletions

File tree

monai/transforms/spatial/array.py

Lines changed: 29 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
from monai.config import USE_COMPILED, DtypeLike
2727
from monai.config.type_definitions import NdarrayOrTensor
2828
from monai.data.box_utils import BoxMode, StandardMode
29-
from monai.data.meta_obj import get_track_meta, set_track_meta
29+
from monai.data.meta_obj import get_track_meta
3030
from monai.data.meta_tensor import MetaTensor
3131
from monai.data.utils import AFFINE_TOL, affine_to_spacing, compute_shape_offset, iter_patch, to_affine_nd, zoom_affine
3232
from monai.networks.layers import AffineTransform, GaussianFilter, grid_pull
@@ -3576,33 +3576,35 @@ def __call__(self, img: torch.Tensor, randomize: bool = True) -> torch.Tensor:
35763576

35773577
if self._do_transform:
35783578
input_shape = img.shape[1:]
3579-
target_shape = tuple(np.round(np.array(input_shape) * self.zoom_factor).astype(np.int_).tolist())
3580-
3581-
resize_tfm_downsample = Resize(
3582-
spatial_size=target_shape, size_mode="all", mode=self.downsample_mode, anti_aliasing=False
3583-
)
3584-
3585-
resize_tfm_upsample = Resize(
3586-
spatial_size=input_shape,
3587-
size_mode="all",
3588-
mode=self.upsample_mode,
3589-
anti_aliasing=False,
3590-
align_corners=self.align_corners,
3579+
# Clamp each axis to at least 1 so F.interpolate never sees a zero-sized dimension.
3580+
target_shape = tuple(max(1, int(np.round(s * self.zoom_factor))) for s in input_shape)
3581+
3582+
# Use F.interpolate directly on a plain tensor to avoid mutating the global
3583+
# set_track_meta flag, which is not thread-safe (see GitHub issue #8409).
3584+
img_t = convert_to_tensor(img, track_meta=False)
3585+
# F.interpolate requires float input and a batch dimension; cast matches
3586+
# the default dtype=float32 that Resize uses internally.
3587+
img_float = img_t.unsqueeze(0).to(dtype=torch.float32)
3588+
3589+
downsample_mode = str(self.downsample_mode)
3590+
upsample_mode = str(self.upsample_mode)
3591+
# align_corners is only valid for linear/bilinear/bicubic/trilinear modes
3592+
_align_corners_modes = {"linear", "bilinear", "bicubic", "trilinear"}
3593+
downsample_align_corners = self.align_corners if downsample_mode in _align_corners_modes else None
3594+
upsample_align_corners = self.align_corners if upsample_mode in _align_corners_modes else None
3595+
3596+
img_downsampled = torch.nn.functional.interpolate(
3597+
img_float, size=target_shape, mode=downsample_mode, align_corners=downsample_align_corners
35913598
)
3592-
# temporarily disable metadata tracking, since we do not want to invert the two Resize functions during
3593-
# post-processing
3594-
original_tack_meta_value = get_track_meta()
3595-
set_track_meta(False)
3596-
3597-
img_downsampled = resize_tfm_downsample(img)
3598-
img_upsampled = resize_tfm_upsample(img_downsampled)
3599-
3600-
# reset metadata tracking to original value
3601-
set_track_meta(original_tack_meta_value)
3602-
3603-
# copy metadata from original image to down-and-upsampled image
3604-
img_upsampled = MetaTensor(img_upsampled)
3605-
img_upsampled.copy_meta_from(img)
3599+
img_upsampled_t = torch.nn.functional.interpolate(
3600+
img_downsampled, size=input_shape, mode=upsample_mode, align_corners=upsample_align_corners
3601+
).squeeze(0)
3602+
3603+
# copy metadata from original image to down-and-upsampled image,
3604+
# respecting the caller's get_track_meta() setting.
3605+
img_upsampled = cast(torch.Tensor, convert_to_tensor(img_upsampled_t, track_meta=get_track_meta()))
3606+
if isinstance(img_upsampled, MetaTensor):
3607+
img_upsampled.copy_meta_from(img)
36063608

36073609
return img_upsampled
36083610

tests/transforms/test_rand_simulate_low_resolution.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,14 @@
1111

1212
from __future__ import annotations
1313

14+
import threading
1415
import unittest
1516

1617
import numpy as np
18+
import torch
1719
from parameterized import parameterized
1820

21+
from monai.data.meta_obj import get_track_meta
1922
from monai.transforms import RandSimulateLowResolution
2023
from tests.test_utils import TEST_NDARRAYS, assert_allclose
2124

@@ -78,6 +81,48 @@ def test_value(self, arguments, image, expected_data):
7881
result = randsimlowres(image)
7982
assert_allclose(result, expected_data, rtol=1e-4, type_test="tensor")
8083

84+
def test_track_meta_global_state_unchanged(self):
85+
# Verify that calling RandSimulateLowResolution does not modify the global
86+
# set_track_meta flag (regression test for GitHub issue #8409).
87+
img = torch.ones(1, 4, 4, 4)
88+
tfm = RandSimulateLowResolution(prob=1.0, zoom_range=(0.5, 0.6))
89+
tfm.set_random_state(seed=0)
90+
91+
original_track_meta = get_track_meta()
92+
tfm(img)
93+
self.assertEqual(get_track_meta(), original_track_meta, "set_track_meta global state was unexpectedly modified")
94+
95+
def test_thread_safety(self):
96+
# Verify that concurrent calls do not corrupt each other's track_meta state
97+
# (regression test for GitHub issue #8409).
98+
# expected_track_meta is captured before threads start so every worker
99+
# checks against the same baseline rather than its own (possibly already
100+
# corrupted) snapshot.
101+
errors = []
102+
expected_track_meta = get_track_meta()
103+
start_barrier = threading.Barrier(8)
104+
105+
def run_transform():
106+
img = torch.ones(1, 4, 4, 4)
107+
tfm = RandSimulateLowResolution(prob=1.0, zoom_range=(0.5, 0.6))
108+
start_barrier.wait() # synchronise so all threads hammer the transform at once
109+
try:
110+
for _ in range(50):
111+
tfm(img)
112+
if get_track_meta() != expected_track_meta:
113+
errors.append(RuntimeError("track_meta state changed in thread"))
114+
break
115+
except Exception as e:
116+
errors.append(e)
117+
118+
threads = [threading.Thread(target=run_transform) for _ in range(8)]
119+
for t in threads:
120+
t.start()
121+
for t in threads:
122+
t.join()
123+
124+
self.assertEqual(errors, [], f"Thread safety errors: {errors}")
125+
81126

82127
if __name__ == "__main__":
83128
unittest.main()

0 commit comments

Comments
 (0)