From 4914e24b1ca48f229dd5cb740867c05492aa8cb2 Mon Sep 17 00:00:00 2001 From: Wenqi Li Date: Tue, 11 Jan 2022 22:17:05 +0000 Subject: [PATCH] update meshgrid Signed-off-by: Wenqi Li --- monai/networks/blocks/crf.py | 3 ++- monai/networks/blocks/warp.py | 3 ++- monai/networks/utils.py | 7 +++++++ monai/transforms/smooth_field/array.py | 8 +++----- monai/transforms/spatial/array.py | 3 ++- tests/test_grid_pull.py | 3 ++- 6 files changed, 18 insertions(+), 9 deletions(-) diff --git a/monai/networks/blocks/crf.py b/monai/networks/blocks/crf.py index ccaef17679..b6382adf5f 100644 --- a/monai/networks/blocks/crf.py +++ b/monai/networks/blocks/crf.py @@ -15,6 +15,7 @@ from torch.nn.functional import softmax from monai.networks.layers.filtering import PHLFilter +from monai.networks.utils import meshgrid_ij __all__ = ["CRF"] @@ -114,6 +115,6 @@ def forward(self, input_tensor: torch.Tensor, reference_tensor: torch.Tensor): # helper methods def _create_coordinate_tensor(tensor): axes = [torch.arange(tensor.size(i)) for i in range(2, tensor.dim())] - grids = torch.meshgrid(axes) + grids = meshgrid_ij(axes) coords = torch.stack(grids).to(device=tensor.device, dtype=tensor.dtype) return torch.stack(tensor.size(0) * [coords], dim=0) diff --git a/monai/networks/blocks/warp.py b/monai/networks/blocks/warp.py index 2cb349f8f0..79fed14834 100644 --- a/monai/networks/blocks/warp.py +++ b/monai/networks/blocks/warp.py @@ -18,6 +18,7 @@ from monai.config.deviceconfig import USE_COMPILED from monai.networks.layers.spatial_transforms import grid_pull +from monai.networks.utils import meshgrid_ij from monai.utils import GridSampleMode, GridSamplePadMode, optional_import _C, _ = optional_import("monai._C") @@ -84,7 +85,7 @@ def __init__(self, mode=GridSampleMode.BILINEAR.value, padding_mode=GridSamplePa @staticmethod def get_reference_grid(ddf: torch.Tensor) -> torch.Tensor: mesh_points = [torch.arange(0, dim) for dim in ddf.shape[2:]] - grid = torch.stack(torch.meshgrid(*mesh_points), dim=0) # (spatial_dims, ...) + grid = torch.stack(meshgrid_ij(*mesh_points), dim=0) # (spatial_dims, ...) grid = torch.stack([grid] * ddf.shape[0], dim=0) # (batch, spatial_dims, ...) grid = grid.to(ddf) return grid diff --git a/monai/networks/utils.py b/monai/networks/utils.py index 5a2bb13fec..a4ca0a6fd5 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -37,6 +37,7 @@ "train_mode", "copy_model_state", "convert_to_torchscript", + "meshgrid_ij", ] @@ -500,3 +501,9 @@ def convert_to_torchscript( torch.testing.assert_allclose(r1, r2, rtol=rtol, atol=atol) return script_module + + +def meshgrid_ij(*tensors): + if pytorch_after(1, 10): + return torch.meshgrid(*tensors, indexing="ij") + return torch.meshgrid(*tensors) diff --git a/monai/transforms/smooth_field/array.py b/monai/transforms/smooth_field/array.py index 31ce76e5b5..356b0d167f 100644 --- a/monai/transforms/smooth_field/array.py +++ b/monai/transforms/smooth_field/array.py @@ -19,11 +19,12 @@ import monai from monai.config.type_definitions import NdarrayOrTensor +from monai.networks.utils import meshgrid_ij from monai.transforms.transform import Randomizable, RandomizableTransform from monai.transforms.utils_pytorch_numpy_unification import moveaxis from monai.utils import GridSampleMode, GridSamplePadMode, InterpolateMode from monai.utils.enums import TransformBackends -from monai.utils.module import look_up_option, pytorch_after +from monai.utils.module import look_up_option from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor __all__ = ["SmoothField", "RandSmoothFieldAdjustContrast", "RandSmoothFieldAdjustIntensity", "RandSmoothDeform"] @@ -404,10 +405,7 @@ def __init__( grid_space = spatial_size if spatial_size is not None else self.sfield.field.shape[2:] grid_ranges = [torch.linspace(-1, 1, d) for d in grid_space] - if pytorch_after(1, 10): - grid = torch.meshgrid(*grid_ranges, indexing="ij") - else: - grid = torch.meshgrid(*grid_ranges) + grid = meshgrid_ij(*grid_ranges) self.grid = torch.stack(grid).unsqueeze(0).to(self.device, self.grid_dtype) diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py index 4e5cac2c85..3e32ddbb18 100644 --- a/monai/transforms/spatial/array.py +++ b/monai/transforms/spatial/array.py @@ -22,6 +22,7 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.data.utils import compute_shape_offset, to_affine_nd, zoom_affine from monai.networks.layers import AffineTransform, GaussianFilter, grid_pull +from monai.networks.utils import meshgrid_ij from monai.transforms.croppad.array import CenterSpatialCrop, Pad from monai.transforms.transform import Randomizable, RandomizableTransform, ThreadUnsafe, Transform from monai.transforms.utils import ( @@ -2103,7 +2104,7 @@ def __call__( ranges = ranges - (dim_size - 1.0) / 2.0 all_ranges.append(ranges) - coords = torch.meshgrid(*all_ranges) + coords = meshgrid_ij(*all_ranges) grid = torch.stack([*coords, torch.ones_like(coords[0])]) return self.resampler(img, grid=grid, mode=mode, padding_mode=padding_mode) # type: ignore diff --git a/tests/test_grid_pull.py b/tests/test_grid_pull.py index b0f19d1a00..561b231498 100644 --- a/tests/test_grid_pull.py +++ b/tests/test_grid_pull.py @@ -16,6 +16,7 @@ from parameterized import parameterized from monai.networks.layers import grid_pull +from monai.networks.utils import meshgrid_ij from monai.utils import optional_import from tests.testing_data.cpp_resample_answers import Expected_1D_GP_bwd, Expected_1D_GP_fwd from tests.utils import skip_if_no_cpp_extension @@ -26,7 +27,7 @@ def make_grid(shape, dtype=None, device=None, requires_grad=True): ranges = [torch.arange(float(s), dtype=dtype, device=device, requires_grad=requires_grad) for s in shape] - grid = torch.stack(torch.meshgrid(*ranges), dim=-1) + grid = torch.stack(meshgrid_ij(*ranges), dim=-1) return grid[None]