rapidsai / dask-cuda

Utilities for Dask and CUDA interactions
https://docs.rapids.ai/api/dask-cuda/stable/
Apache License 2.0
292 stars 93 forks source link

localcudacluster with onnxruntime inference #1399

Open kanglcn opened 1 month ago

kanglcn commented 1 month ago

Hi,

I am working on using a trained deep-learning model for image denoising. The model is saved in onnx format and I successfully deployed this model with onnxruntime. The workflow is:

  1. convert numpy array to cupy array
  2. do some preprocessing on cupy array
  3. create the onnxruntime session with gpu support
  4. run the model inference with input and out binding to cupy array
  5. do some afterprocessing on cupy array
  6. convert cupy array back to numpy array

Since I have many images to denoise and a signal-node-multi-gpu machine, I wrap the above workflow to one function and I want to use dask-cuda to automatically distribute these tasks. However, the worker always died unreasonably.

I did one test on other cupy-only processing workflow and it works. But with onnxruntime, it never works. I would appreciate it if anybody can help!

Thanks!