From a7efe450b43cdb55acad65d7d4320044e788a39b Mon Sep 17 00:00:00 2001 From: Nic Ma Date: Mon, 30 Aug 2021 12:26:38 +0800 Subject: [PATCH 1/2] [DLMED] add args to ThreadDataLoader Signed-off-by: Nic Ma --- monai/data/thread_buffer.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/monai/data/thread_buffer.py b/monai/data/thread_buffer.py index 2901335bd5..df33f6ed84 100644 --- a/monai/data/thread_buffer.py +++ b/monai/data/thread_buffer.py @@ -33,7 +33,7 @@ class ThreadBuffer: timeout: Time to wait for an item from the buffer, or to wait while the buffer is full when adding items """ - def __init__(self, src, buffer_size=1, timeout=0.01): + def __init__(self, src, buffer_size: int = 1, timeout: float = 0.01): self.src = src self.buffer_size = buffer_size self.timeout = timeout @@ -82,11 +82,27 @@ class ThreadDataLoader(DataLoader): Subclass of `DataLoader` using a `ThreadBuffer` object to implement `__iter__` method asynchronously. This will iterate over data from the loader as expected however the data is generated on a separate thread. Use this class where a `DataLoader` instance is required and not just an iterable object. + + Args: + dataset: input dataset. + buffer_size: number of items to buffer from the data source. + buffer_timeout: time to wait for an item from the buffer, or to wait while the buffer is full when adding items. + num_workers: number of the multi-prcessing workers in PyTorch DataLoader. + """ - def __init__(self, dataset: Dataset, num_workers: int = 0, **kwargs): + def __init__( + self, + dataset: Dataset, + buffer_size: int = 1, + buffer_timeout: float = 0.01, + num_workers: int = 0, + **kwargs, + ): super().__init__(dataset, num_workers, **kwargs) + self.buffer_size = buffer_size + self.buffer_timeout = buffer_timeout def __iter__(self): - buffer = ThreadBuffer(super().__iter__()) + buffer = ThreadBuffer(src=super().__iter__(), buffer_size=self.buffer_size, timeout=self.buffer_timeout) yield from buffer From 84b066be60a0dd96bf857b1c45cc8701a50c4944 Mon Sep 17 00:00:00 2001 From: Nic Ma Date: Mon, 30 Aug 2021 12:32:16 +0800 Subject: [PATCH 2/2] [DLMED] fix flake8 Signed-off-by: Nic Ma --- monai/data/thread_buffer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/data/thread_buffer.py b/monai/data/thread_buffer.py index df33f6ed84..da5847465e 100644 --- a/monai/data/thread_buffer.py +++ b/monai/data/thread_buffer.py @@ -37,7 +37,7 @@ def __init__(self, src, buffer_size: int = 1, timeout: float = 0.01): self.src = src self.buffer_size = buffer_size self.timeout = timeout - self.buffer = Queue(self.buffer_size) + self.buffer: Queue = Queue(self.buffer_size) self.gen_thread = None self.is_running = False