From e5d790786f532c6d53e288b7b1a07e947fe41f36 Mon Sep 17 00:00:00 2001 From: Yiheng Wang Date: Tue, 17 Dec 2024 09:44:44 +0000 Subject: [PATCH] update Signed-off-by: Yiheng Wang --- monai/data/image_reader.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/monai/data/image_reader.py b/monai/data/image_reader.py index 6598fc3829..86702f087a 100644 --- a/monai/data/image_reader.py +++ b/monai/data/image_reader.py @@ -16,6 +16,7 @@ import io import os import re +import tempfile import warnings from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator, Sequence @@ -909,9 +910,30 @@ def __init__( ) to_gpu = False + if to_gpu: + self.warmup_kvikio() + self.to_gpu = to_gpu self.kwargs = kwargs + def warmup_kvikio(self): + """ + Warm up the Kvikio library to initialize the internal buffers, cuFile, GDS, etc. + This can accelerate the data loading process when `to_gpu` is set to True. + """ + if has_cp and has_kvikio: + print("warm up") + a = cp.arange(100) + with tempfile.NamedTemporaryFile() as tmp_file: + tmp_file_name = tmp_file.name + f = kvikio.CuFile(tmp_file_name, "w") + f.write(a) + f.close() + + b = cp.empty_like(a) + f = kvikio.CuFile(tmp_file_name, "r") + f.read(b) + def verify_suffix(self, filename: Sequence[PathLike] | PathLike) -> bool: """ Verify whether the specified file or files format is supported by Nibabel reader. @@ -961,6 +983,9 @@ def get_data(self, img) -> tuple[np.ndarray, dict]: img: a Nibabel image object loaded from an image file or a list of Nibabel image objects. """ + # TODO: the actual type is list[np.ndarray | cp.ndarray] + # should figure out how to define correct types without having cupy not found error + # https://github.com/Project-MONAI/MONAI/pull/8188#discussion_r1886645918 img_array: list[np.ndarray] = [] compatible_meta: dict = {}