[FEA] Provide a way/example to use cuCIM with Dask for DataLoading(Pytorch's Dataloader-like API)

[gigony]

Is your feature request related to a problem? Please describe.

PyTorch's DataLoader class is used in many DeepLearning training applications to load training data and pre-process the data, before feeding to AI model.

Since PyTorch's DataLoader is running in multi-processes, it is hard to use cuCIM's scikit-image APIs (which makes use of CUDA) in the pre-transformations of the DataLoader due to CUDA context issues.

It would be nice to provide a way/example to use cuCIM with DeepLearning Frameworks such as PyTorch.

Describe the solution you'd like

PyTorch's DataLoader works like this. It would be nice if we have a PyTorch's DataLoader-like utility class in Dask that mimics Pytorch's DataLoader behavior but implemented with Dask (dask-cuda) for the parallelization of data loading (so providing a generator/iterator that gives a batch of processed image data).

Describe alternatives you've considered

• To use cuCIM in the training pipeline, we currently move GPU-accelerated pre-transforms from PyTorch DataLoader's transformation(using Compose) to the main thread (place GPU-based batch pre-transformation right before feeding to the AI model, and right after getting CPU-loaded/pre-transformed training data by DataLoader.), to avoid CUDA context issues.
• It would be good if we also provide an example with that approach.

Additional context

Relevant information regarding CuPy+PyTorch.

• Interoperability with PyTorch memory pool cupy/cupy#2762
• https://pytorch-pfn-extras.readthedocs.io/en/latest/user_guide/cuda.html
• Force libtorch to use CUDA context pytorch/pytorch#31565
• Support for Per Thread Default Stream (PTDS) cupy/cupy#4322

With Numba to get cuda context.

• https://github.com/numba/numba/blob/1fc53ecb6183f441498c0e082e748e6792e29791/numba/cuda/cudadrv/devices.py#L130

[NV-jpt]

I would be happy to take a look at this!

[rjzamora]

NVTabular's pytorch dataloader (which is built on Dask) may be agood reference here. I'll be happy to help advise on this work and clarify what NVTabular is (and isn't) doing.

[NV-jpt]

Thank you, @rjzamora !

[NV-jpt]

Unfortunately, I have come across some issues with the cuCIM codebase that are blocking this Dask-based DataLoading solution.

For example, while the array transformation in the cuCIM transform image_rotate_90() should be compatible with a Dask Array input, there is an explicit Type-Check here that throws a TypeError whenever trying to apply cuCIM transforms to a Dask Array.

In order to allow Dask to schedule cuCIM operations - we will likely want to make this "check" more of a duck-type check, that checks for the necessary API interface, as Dask Arrays should be able to pass those checks.