Repository: Project-MONAI/MONAI

Issue:
`GridSplit` torch vs numpy performance 
`GridSplit` has two separate implementations for `torch.Tensor` and `numpy.ndarray`. Although the torch implementation using `unfold` is specific to torch, both should be able to use the numpy implementation using `as_strided`. Thus, it would be nice to compare the performance of these two implementation and if `as_strided` is better, we can unify the implementation. Other factor to consider would be ability to expand to tensors with three spatial dimensions.



Discussion:
Presumably even if `np.as_strided` is faster, you would still have to use `torch.unfold` if a `torch.Tensor` was not on the CPU?
> Presumably even if `np.as_strided` is faster, you would still have to use `torch.unfold` if a `torch.Tensor` was not on the CPU?

Yes, in the current implementation we use `unfold` for any `torch.Tensor` or CPU or GPU. But if `torch.as_strided` is faster than `torch.unfold`, then we can use [`torch.as_strided`](https://pytorch.org/docs/stable/generated/torch.as_strided.html#torch.as_strided) instead. I am not suggesting to convert from torch to numpy but use the same API in torch and numpy to simplify the code if it also makes it faster.

https://github.com/Project-MONAI/MONAI/blob/d7767cd31e89fa9490ce958e97b253e27c3b62f8/monai/transforms/spatial/array.py#L2553-L2571
Oh I see, that definitely makes sense. Keen to see the results from the profiling!

Resolving patch:
diff --git a/monai/transforms/spatial/array.py b/monai/transforms/spatial/array.py
--- a/monai/transforms/spatial/array.py
+++ b/monai/transforms/spatial/array.py
@@ -15,11 +15,10 @@
 import warnings
 from copy import deepcopy
 from enum import Enum
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
-from numpy.lib.stride_tricks import as_strided
 
 from monai.config import USE_COMPILED, DtypeLike
 from monai.config.type_definitions import NdarrayOrTensor
@@ -2970,31 +2969,31 @@ def __call__(
 
         split_size, steps = self._get_params(image.shape[1:], input_size)
         patches: List[NdarrayOrTensor]
+        as_strided_func: Callable
         if isinstance(image, torch.Tensor):
-            unfolded_image = (
-                image.unfold(1, split_size[0], steps[0])
-                .unfold(2, split_size[1], steps[1])
-                .flatten(1, 2)
-                .transpose(0, 1)
-            )
-            # Make a list of contiguous patches
-            patches = [p.contiguous() for p in unfolded_image]
+            as_strided_func = torch.as_strided
+            c_stride, x_stride, y_stride = image.stride()  # type: ignore
         elif isinstance(image, np.ndarray):
-            x_step, y_step = steps
+            as_strided_func = np.lib.stride_tricks.as_strided
             c_stride, x_stride, y_stride = image.strides
-            n_channels = image.shape[0]
-            strided_image = as_strided(
-                image,
-                shape=(*self.grid, n_channels, split_size[0], split_size[1]),
-                strides=(x_stride * x_step, y_stride * y_step, c_stride, x_stride, y_stride),
-            )
-            # Flatten the first two dimensions
-            strided_image = strided_image.reshape(-1, *strided_image.shape[2:])
-            # Make a list of contiguous patches
-            patches = [np.ascontiguousarray(p) for p in strided_image]
         else:
             raise ValueError(f"Input type [{type(image)}] is not supported.")
 
+        x_step, y_step = steps
+        n_channels = image.shape[0]
+        strided_image = as_strided_func(
+            image,
+            (*self.grid, n_channels, split_size[0], split_size[1]),
+            (x_stride * x_step, y_stride * y_step, c_stride, x_stride, y_stride),
+        )
+        # Flatten the first two dimensions
+        strided_image = strided_image.reshape(-1, *strided_image.shape[2:])
+        # Make a list of contiguous patches
+        if isinstance(image, torch.Tensor):
+            patches = [p.contiguous() for p in strided_image]
+        elif isinstance(image, np.ndarray):
+            patches = [np.ascontiguousarray(p) for p in strided_image]
+
         return patches
 
     def _get_params(


Test patch: