TORCH backend for RandHistogramShift transform

[dyollb]

Is your feature request related to a problem? Please describe.
Currently the only transform in our augmentation pipeline that doesn't have a TORCH backend is the RandHistogramShiftd.
I have not profiled the speedup and GPU memory for the augmentation transforms, but I guess if we could do most of it on the GPU it should improve performance.

Describe the solution you'd like
Implement TORCH backend for RandHistogramShift.

Describe alternatives (solution) you've considered
I would propose to

• not change RandHistogramShift.randomize, i.e. control points remain np.ndarray
• cast the control points to the input NdarrayOrTensor type
• implement a function with the functionality/API of np.interp (used in __call__) that supports NdarrayOrTensor

An example implementation for np.interp-like function using torch.Tensor is mentioned here: pytorch/pytorch#50334

If you agree to this strategy I will submit a PR.

[dyollb]

How about something like this?:

class RandHistogramShift(RandomizableTransform):

    backend = [TransformBackends.TORCH, TransformBackends.NUMPY]

    def __init__(self, num_control_points: Union[Tuple[int, int], int] = 10, prob: float = 0.1) -> None:
        ...

    def randomize(self, data: Optional[Any] = None) -> None:
        ...

    def interp(x: NdarrayOrTensor, xp: NdarrayOrTensor, fp: NdarrayOrTensor) -> NdarrayOrTensor:

        m = (fp[1:]- fp[:-1]) / (xp[1:]- xp[:-1])
        b = fp[:-1] - (m * xp[:-1])

        indices = (x.reshape(-1, 1) >= xp.reshape(1, -1)).sum(1) - 1
        indices = indices.clip(0, len(m) - 1)

        return (m[indices] * x.reshape(-1) + b[indices]).reshape( x.shape)

    def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor:
        if randomize:
            self.randomize()

        if not self._do_transform:
            return img

        if self.reference_control_points is None or self.floating_control_points is None:
            raise RuntimeError("please call the `randomize()` function first.")
        
        xp = convert_to_dst_type(self.reference_control_points, dst=img)
        yp = convert_to_dst_type(self.floating_control_points, dst=img)
        img_min, img_max = img.min(), img.max()
        reference_control_points_scaled = xp * (img_max - img_min) + img_min
        floating_control_points_scaled = yp * (img_max - img_min) + img_min
        img = self.interp(img, reference_control_points_scaled, floating_control_points_scaled)
        return img

[Nic-Ma]

Hi @dyollb ,

Thanks for your idea.
I think it would be nice to support torch backend for all the MONAI transforms.
Please feel free to contribute a PR for it.

Thanks in advance.

[dyollb]

Sure!

Btw. here are some benchmarks for the interp like function (np.ndarray, tensor, cpu/gpu):
https://gist.github.com/dyollb/e4804f8d22d837b27778450e7f9ebce6