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:
gml16: Building on your suggestions @Nic-Ma, I implemented the CacheDataset with 2 modes for now (threads and processes) and timed them. Threads is much faster than processes and if you agree I'll remove the multi-processes mode.

I have included the code I used to time both threading and processes in `tests/test_cachedataset_parallel.py`. I have an Intel i7 CPU which does not compare to #309 use case.

Here are the timing results:

```
python tests/test_cachedataset_parallel.py
Load and cache transformed data...
100/100 [==============================]
0 workers in mode thread took 11.77751111984253s to cache 100 items.
Load and cache transformed data...
100/100 [==============================]
0 workers in mode process took 11.39584469795227s to cache 100 items.
Load and cache transformed data...
100/100 [==============================]
1 workers in mode thread took 11.767529726028442s to cache 100 items.
Load and cache transformed data...
50/50 [==============================]
1 workers in mode process took 148.93750596046448s to cache 50 items.
Load and cache transformed data...
100/100 [==============================]
4 workers in mode thread took 8.87144923210144s to cache 100 items.
Load and cache transformed data...
50/50 [==============================]
4 workers in mode process took 138.71821284294128s to cache 50 items.
```
Nic-Ma: Thanks for your update.
I tried to run your program in the spleen segmentation notebook
The multi-thread version takes 95.75s to cache data, but the multi-process version still can't show the process bar. Could you please help double-check the process bar?
Thanks.
Nic-Ma: BTW, there is another problem that:
You call `self._cache.append(item)` in multi-thread or multi-process which can't guarantee the original data item order and equal to an unwanted `shuffle`.
It will affect the deterministic training and evaluation results, and can't save model results in order.
Maybe you can try to add **ID** and change to **dict** type.
Thanks.
gml16: Thank you for your inputs. 
I did not think about the concurrent shuffle, I will update this soon.
Regarding the process bar, I have also tried the code on Jupyter notebook and it doesn't seem to work in process mode. However, it does work when I'm running the code in a `.py` file.

![image](https://user-images.githubusercontent.com/7902440/80521942-b9d9c800-8983-11ea-9a4b-732499e1b055.png)
 
I'm wondering if this is because Jupyter notebooks behave differently with multiprocessing. Will look further into it.
Nic-Ma: Hi @gml16 ,

I think the `multi-thread` method in your PR is almost good enough, just need to fix multi-process process bar in notebook and compare the speed with `multi-process` to double-confirm.
Thanks.
gml16: I have just pushed an update which should resolve most of the points above.
The process bar in Jupyter in process mode is displayed in the terminal instead of the output cell. I am not sure if there is a way around this as the multiprocessing package does not always work in interactive interpreters:

![image](https://user-images.githubusercontent.com/7902440/80740381-e589b900-8b0f-11ea-8115-37726eb739c2.png)
(from https://docs.python.org/3.8/library/multiprocessing.html) 

In thread mode the process bar is shown directly in the output cell in a Jupyter notebook.

Speed wise here is what I got:

![image](https://user-images.githubusercontent.com/7902440/80741065-f981ea80-8b10-11ea-9648-aece0638deb7.png)

The third cell is in the case where the code is not parallelised. And as mentioned above, the process bar for the second cell shows in the terminal.
ericspod: Multiprocessing has a number of odd gotchas due to the way processes are spawned in different OSes and the information that needs to be shared to operate in the expected way. Python kernels in Jupyter notebooks communicate with the front end through ZeroMQ, I'm not sure subprocesses get the correct handles to do that without being explicitly coded to do so. Also they will write to stdout which isn't being captured in the same way as it is for the parent process. One solution in multiprocessor mode would be to use [imap_unordered](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool.imap_unordered) and iterate over the results coming from it, incrementing the counter as you go:

```python
from multiprocessing import Pool
with Pool() as p:
    count=0
    for _ in p.imap_unordered(do_work,range(10),1):
        counter+=1
        print(counter) # do the progress bar print here instead
```

The parent process is responsible for printing the output, not the child processes. The parent is the kernel attached to the notebook so its output will be in the cell as expected. This means you don't need to share the counter using a `Value` object (which has its own lock accessed through [get_lock()](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Value) incidentally), and this pattern can be used for the thread version as well.
Nic-Ma: Hi @gml16  and @ericspod ,

Thanks for your investigation and testing!
According to the test result, I think we can only keep the `multi-threads` version and merge now.
Of course, we can add `multi-process` version if find some unknown benefits in the future?
Do you agree?
Thanks.
gml16: Thank you @ericspod for the detailed explanation. @Nic-Ma I've removed the multiprocess mode for now and it should be good for merging. Thanks for all the help.

Patch:
diff --git a/monai/data/dataset.py b/monai/data/dataset.py
index 9d4b3b92..80866031 100644
--- a/monai/data/dataset.py
+++ b/monai/data/dataset.py
@@ -13,6 +13,9 @@ import sys
 
 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 @@ class CacheDataset(Dataset):
                 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 @@ class CacheDataset(Dataset):
         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
index 89f0111f..14e32e12 100644
--- a/tests/test_cachedataset.py
+++ b/tests/test_cachedataset.py
@@ -23,7 +23,6 @@ TEST_CASE_1 = [
     (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
index 00000000..aa0b4a7a
--- /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()