Repository: Project-MONAI/MONAI

Issue:
CacheDataset would benefit greatly from multithreaded processing
**Is your feature request related to a problem? Please describe.**
I have access to huge amounts of medical imaging data (20,000 3D scans), and very large many core computer (88 hyperthreaded CPU's).

The pre-preocessing step to generate the initial cache is currently single threaded, so building up the initial cache is taking much longer than expected.

**Describe the solution you'd like**

I think that supporting multi-threading for applying the initial transorms would be a huge benefit to testing.


```python
CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0)

##
CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=88)
##---------------------------------------------------------------------^^^^^^^^^^^^^^^^^^
```
**Additional context**
After looking at the DataLoader interface for implementing  `num_workers`, I realized that this may not be as trivial of a task as I expected.


Discussion:
Hi @hjmjohnson ,

Thanks very much for raising this topic!
Actually, we already discussed it internally before, it's a very important enhancement for our CacheDataset mechanism. Due to schedule, I didn't add multi-thread support for v0.1.0 release.
Would you like to submit a PR for `num_threads` feature if you have time?
Of course, if you are not available, I can also try to implement it several weeks later when I finished existing tasks.
Thanks in advance.

Resolving patch:
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -13,6 +13,9 @@
 
 import torch
 
+from multiprocessing.pool import ThreadPool
+import threading
+
 from monai.transforms.compose import Compose, Randomizable
 from monai.transforms.utils import apply_transform
 from monai.utils import process_bar
@@ -84,7 +87,7 @@ class CacheDataset(Dataset):
     and the outcome not cached.
     """
 
-    def __init__(self, data, transform, cache_num=sys.maxsize, cache_rate=1.0):
+    def __init__(self, data, transform, cache_num=sys.maxsize, cache_rate=1.0, num_workers=0):
         """
         Args:
             data (Iterable): input data to load and transform to generate dataset for model.
@@ -93,22 +96,39 @@ def __init__(self, data, transform, cache_num=sys.maxsize, cache_rate=1.0):
                 will take the minimum of (cache_num, data_length x cache_rate, data_length).
             cache_rate (float): percentage of cached data in total, default is 1.0 (cache all).
                 will take the minimum of (cache_num, data_length x cache_rate, data_length).
+            num_workers (int): the number of worker processes to use.
+                If 0 a single thread will be used. Default is 0.
         """
         if not isinstance(transform, Compose):
             transform = Compose(transform)
         super().__init__(data, transform)
         self.cache_num = min(cache_num, int(len(self) * cache_rate), len(self))
-        self._cache = list()
+        self._cache = [None] * self.cache_num
         print('Load and cache transformed data...')
-        for i in range(self.cache_num):
-            process_bar(i + 1, self.cache_num)
-            item = data[i]
-            for _transform in transform.transforms:
-                # execute all the deterministic transforms before the first random transform
-                if isinstance(_transform, Randomizable):
-                    break
-                item = apply_transform(_transform, item)
-            self._cache.append(item)
+        if num_workers > 0:
+            self._item_processed = 0
+            self._thread_lock = threading.Lock()
+            with ThreadPool(num_workers) as p:
+                p.map(self._load_cache_item_thread, [(i, data[i], transform.transforms) for i in range(self.cache_num)])
+        else:
+            for i in range(self.cache_num):
+                self._cache[i] = self._load_cache_item(data[i], transform.transforms)
+                process_bar(i + 1, self.cache_num)
+
+    def _load_cache_item(self, item, transforms):
+        for _transform in transforms:
+            # execute all the deterministic transforms before the first random transform
+            if isinstance(_transform, Randomizable):
+                break
+            item = apply_transform(_transform, item)
+        return item
+
+    def _load_cache_item_thread(self, args):
+        i, item, transforms = args
+        self._cache[i] = self._load_cache_item(item, transforms)
+        with self._thread_lock:
+            self._item_processed += 1
+            process_bar(self._item_processed, self.cache_num)
 
     def __getitem__(self, index):
         if index < self.cache_num:
@@ -124,5 +144,4 @@ def __getitem__(self, index):
         else:
             # no cache for this data, execute all the transforms directly
             data = super(CacheDataset, self).__getitem__(index)
-
         return data


Test patch:
diff --git a/tests/test_cachedataset.py b/tests/test_cachedataset.py
--- a/tests/test_cachedataset.py
+++ b/tests/test_cachedataset.py
@@ -23,7 +23,6 @@
     (128, 128, 128)
 ]
 
-
 class TestCacheDataset(unittest.TestCase):
 
     @parameterized.expand([TEST_CASE_1])
diff --git a/tests/test_cachedataset_parallel.py b/tests/test_cachedataset_parallel.py
new file mode 100644
--- /dev/null
+++ b/tests/test_cachedataset_parallel.py
@@ -0,0 +1,59 @@
+# Copyright 2020 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import shutil
+import numpy as np
+import tempfile
+import nibabel as nib
+from parameterized import parameterized
+from monai.data import CacheDataset
+from monai.transforms import Compose, LoadNiftid
+
+TEST_CASE_1 = [
+    0, 100
+]
+
+TEST_CASE_2 = [
+    4, 100
+]
+
+
+class TestCacheDatasetParallel(unittest.TestCase):
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_shape(self, num_workers, dataset_size):
+        test_image = nib.Nifti1Image(np.random.randint(0, 2, size=[128, 128, 128]), np.eye(4))
+        tempdir = tempfile.mkdtemp()
+        nib.save(test_image, os.path.join(tempdir, 'test_image1.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_label1.nii.gz'))
+        nib.save(test_image, os.path.join(tempdir, 'test_extra1.nii.gz'))
+        test_data = [
+            {
+                'image': os.path.join(tempdir, 'test_image1.nii.gz'),
+                'label': os.path.join(tempdir, 'test_label1.nii.gz'),
+                'extra': os.path.join(tempdir, 'test_extra1.nii.gz')
+            }
+        ] * dataset_size
+        dataset = CacheDataset(
+            data=test_data,
+            transform=Compose([LoadNiftid(keys=['image', 'label', 'extra'])]),
+            cache_rate=1,
+            num_workers=num_workers
+        )
+        shutil.rmtree(tempdir)
+        self.assertEqual(len(dataset._cache), dataset.cache_num)
+        for i in range(dataset.cache_num):
+            self.assertIsNotNone(dataset._cache[i])
+
+if __name__ == '__main__':
+    unittest.main()