Repository: Project-MONAI/MONAI

Issue:
Add partition dataset API
**Is your feature request related to a problem? Please describe.**
In order to partition dataset for training, evaluation, cross validation, even distributed training, we should develop a general `partition_dataset` API.
Refer to what we have in the tutorial:
https://github.com/Project-MONAI/tutorials/blob/master/acceleration/distributed_training/unet_training_smartcache.py#L80


Discussion:

Resolving patch:
diff --git a/monai/data/utils.py b/monai/data/utils.py
--- a/monai/data/utils.py
+++ b/monai/data/utils.py
@@ -12,12 +12,14 @@
 import math
 import os
 import warnings
+from collections import defaultdict
 from itertools import product, starmap
 from pathlib import PurePath
 from typing import Dict, Generator, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
+from torch.utils.data import DistributedSampler as _TorchDistributedSampler
 from torch.utils.data._utils.collate import default_collate
 
 from monai.networks.layers.simplelayers import GaussianFilter
@@ -548,3 +550,203 @@ def is_supported_format(filename: Union[Sequence[str], str], suffixes: Sequence[
             return False
 
     return True
+
+
+def partition_dataset(
+    data: Sequence,
+    ratios: Optional[Sequence[float]] = None,
+    num_partitions: Optional[int] = None,
+    shuffle: bool = False,
+    seed: int = 0,
+    drop_last: bool = False,
+    even_divisible: bool = False,
+):
+    """
+    Split the dataset into N partitions. It can support shuffle based on specified random seed.
+    Will return a set of datasets, every dataset contains 1 partion of original dataset.
+    And it can split the dataset based on specified ratios or evenly split into `num_partitions`.
+    Refer to: https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py.
+
+    Args:
+        data: input dataset to split, expect a list of data.
+        ratios: a list of ratio number to split the dataset, like [8, 1, 1].
+        num_partitions: expected number of the partitions to evenly split, only works when no `ratios`.
+        shuffle: whether to shuffle the original dataset before splitting.
+        seed: random seed to shuffle the dataset, only works when `shuffle` is True.
+        drop_last: only works when `even_divisible` is False and no ratios specified.
+            if True, will drop the tail of the data to make it evenly divisible across partitions.
+            if False, will add extra indices to make the data evenly divisible across partitions.
+        even_divisible: if True, guarantee every partition has same length.
+
+    Examples:
+        data: [1, 2, 3, 4, 5]
+        (1) ratios: [0.6, 0.2, 0.2], shuffle=False, output: [[1, 2, 3], [4], [5]]
+        num_partitions=2, shuffle=False
+        (2) even_divisible=True, drop_last=True, output: [[1, 3], [2, 4]]
+        (3) even_divisible=True, drop_last=False, output: [[1, 3, 5], [2, 4, 1]]
+        (4) even_divisible=False, drop_last=False, output: [[1, 3, 5], [2, 4]]
+
+    """
+    data_len = len(data)
+    datasets = list()
+
+    indices = list(range(data_len))
+    if shuffle:
+        # deterministically shuffle based on fixed seed for every process
+        np.random.seed(seed)
+        np.random.shuffle(indices)
+
+    if ratios is not None:
+        start_idx = next_idx = 0
+        rsum = sum(ratios)
+        for r in ratios:
+            start_idx = next_idx
+            next_idx = min(start_idx + int(r / rsum * data_len + 0.5), data_len)
+            datasets.append([data[i] for i in indices[start_idx:next_idx]])
+    else:
+        if num_partitions is None:
+            raise ValueError("must specify number of partitions.")
+        # evenly split the data without ratios
+        if not even_divisible and drop_last:
+            raise RuntimeError("drop_last only works when even_divisible is True.")
+        if data_len < num_partitions:
+            raise RuntimeError(f"there is no enough data to be splitted for {num_partitions} partitions.")
+
+        if drop_last and data_len % num_partitions != 0:
+            # split to nearest available length that is evenly divisible
+            num_samples = math.ceil((data_len - num_partitions) / num_partitions)
+        else:
+            num_samples = math.ceil(data_len / num_partitions)
+        # use original data length if not even divisible
+        total_size = num_samples * num_partitions if even_divisible else data_len
+
+        if not drop_last and total_size - data_len > 0:
+            # add extra samples to make it evenly divisible
+            indices += indices[: (total_size - data_len)]
+        else:
+            # remove tail of data to make it evenly divisible
+            indices = indices[:total_size]
+
+        for i in range(num_partitions):
+            _indices = indices[i:total_size:num_partitions]
+            datasets.append([data[j] for j in _indices])
+
+    return datasets
+
+
+def partition_dataset_classes(
+    data: Sequence,
+    classes: Sequence[int],
+    ratios: Optional[Sequence[float]],
+    num_partitions: Optional[int],
+    shuffle: bool = False,
+    seed: int = 0,
+    drop_last: bool = False,
+    even_divisible: bool = False,
+):
+    """
+    Split the dataset into N partitions based on the given class labels.
+    It can make sure the same ratio of classes in every partition.
+    Others are same as :py:class:`monai.data.partition_dataset`.
+
+    Args:
+        data: input dataset to split, expect a list of data.
+        classes: a list of labels to help split the data, the length must match the length of data.
+        ratios: a list of ratio number to split the dataset, like [8, 1, 1].
+        num_partitions: expected number of the partitions to evenly split, only works when no `ratios`.
+        shuffle: whether to shuffle the original dataset before splitting.
+        seed: random seed to shuffle the dataset, only works when `shuffle` is True.
+        drop_last: only works when `even_divisible` is False and no ratios specified.
+            if True, will drop the tail of the data to make it evenly divisible across partitions.
+            if False, will add extra indices to make the data evenly divisible across partitions.
+        even_divisible: if True, guarantee every partition has same length.
+
+    Examples:
+        data: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
+        classes: [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3]
+        shuffle: False, ratios: [2, 1]
+        output: [[2, 8, 4, 1, 3, 6, 5, 11, 12], [10, 13, 7, 9, 14]]
+
+    """
+    data_len = len(data)
+    datasets = list()
+
+    if classes is not None:
+        if len(classes) != data_len:
+            raise ValueError("length of classes must match the dataset length.")
+        class_indices = defaultdict(list)
+        for i, c in enumerate(classes):
+            class_indices[c].append(i)
+
+        class_partition_indices: List[Sequence] = list()
+        for _, per_class_indices in sorted(class_indices.items()):
+            per_class_partition_indices = partition_dataset(
+                data=per_class_indices,
+                ratios=ratios,
+                num_partitions=num_partitions,
+                shuffle=shuffle,
+                seed=seed,
+                drop_last=drop_last,
+                even_divisible=even_divisible,
+            )
+            if len(class_partition_indices) == 0:
+                class_partition_indices = per_class_partition_indices
+            else:
+                for part, data_indices in zip(class_partition_indices, per_class_partition_indices):
+                    part += data_indices
+
+        for indices in class_partition_indices:
+            if shuffle:
+                np.random.seed(seed)
+                np.random.shuffle(indices)
+            datasets.append([data[j] for j in indices])
+
+    return datasets
+
+
+def generate_cross_validation_fold(partitions: Sequence[List], fold_idx: int):
+    """
+    Generate cross validation fold based on data partitions and specified fold index.
+    Will select the partition with fold index as validaiton dataset and concat others as train dataset.
+    For example, `partitions`: [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] and `fold_idx`: 2
+    The the output will be [1, 2, 3, 4, 7, 8, 9, 10], [5, 6]
+
+    """
+    if fold_idx >= len(partitions):
+        raise ValueError(f"fold_idx: {fold_idx} is bigger than number of partitions.")
+
+    val_list = partitions[fold_idx]
+    train_list = list()
+    for i, data in enumerate(partitions):
+        if i != fold_idx:
+            train_list += data
+
+    return train_list, val_list
+
+
+class DistributedSampler(_TorchDistributedSampler):
+    """
+    Enhance PyTorch DistributedSampler to support non-evenly divisible sampling.
+
+    Args:
+        even_divisible: if False, different ranks can have different data length.
+        for example, input data: [1, 2, 3, 4, 5], rank 0: [1, 3, 5], rank 1: [2, 4].
+
+    More information about DistributedSampler, please check:
+    https://github.com/pytorch/pytorch/blob/master/torch/utils/data/distributed.py
+
+    """
+
+    def __init__(self, even_divisible: bool = True, **kwargs):
+        self.total_size: int = 0
+        self.rank: int = 0
+        self.num_samples: int = 0
+        self.num_replicas: int = 0
+        super().__init__(**kwargs)
+
+        if not even_divisible:
+            data_len = len(kwargs["dataset"])
+            extra_size = self.total_size - data_len
+            if self.rank + extra_size >= self.num_replicas:
+                self.num_samples -= 1
+            self.total_size = data_len


Test patch:
diff --git a/tests/test_distributed_sampler.py b/tests/test_distributed_sampler.py
new file mode 100644
--- /dev/null
+++ b/tests/test_distributed_sampler.py
@@ -0,0 +1,47 @@
+# 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 numpy as np
+import torch
+import torch.distributed as dist
+
+from monai.data import DistributedSampler
+
+
+def test(expected, **kwargs):
+    dist.init_process_group(backend="nccl", init_method="env://")
+
+    torch.cuda.set_device(dist.get_rank())
+    data = [1, 2, 3, 4, 5]
+    sampler = DistributedSampler(dataset=data, **kwargs)
+    samples = np.array([data[i] for i in list(sampler)])
+    if dist.get_rank() == 0:
+        np.testing.assert_allclose(samples, np.array(expected[0]))
+
+    if dist.get_rank() == 1:
+        np.testing.assert_allclose(samples, np.array(expected[1]))
+
+    dist.destroy_process_group()
+
+
+def main():
+    test(shuffle=False, expected=[[1, 3, 5], [2, 4, 1]])
+    test(shuffle=False, even_divisible=False, expected=[[1, 3, 5], [2, 4]])
+
+
+# suppose to execute on 2 rank processes
+# python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_PER_NODE
+#        --nnodes=NUM_NODES --node_rank=INDEX_CURRENT_NODE
+#        --master_addr="localhost" --master_port=1234
+#        test_distributed_sampler.py
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_generate_cross_validation_fold.py b/tests/test_generate_cross_validation_fold.py
new file mode 100644
--- /dev/null
+++ b/tests/test_generate_cross_validation_fold.py
@@ -0,0 +1,53 @@
+# 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
+
+from parameterized import parameterized
+
+from monai.data import generate_cross_validation_fold, partition_dataset
+
+TEST_CASE_1 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "num_partitions": 5,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[1, 3, 4, 5], [2]],
+]
+
+TEST_CASE_2 = [
+    {
+        "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+        "num_partitions": 10,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[1, 3, 4, 5, 6, 7, 8, 9, 10], [2]],
+]
+
+
+class TestGenerateCrossValidationFold(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_value(self, input_param, result):
+        partitions = partition_dataset(**input_param)
+        train, val = generate_cross_validation_fold(partitions, fold_idx=1)
+        self.assertListEqual(train, result[0])
+        self.assertListEqual(val, result[1])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_partition_dataset.py b/tests/test_partition_dataset.py
new file mode 100644
--- /dev/null
+++ b/tests/test_partition_dataset.py
@@ -0,0 +1,136 @@
+# 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
+
+from parameterized import parameterized
+
+from monai.data import partition_dataset
+
+TEST_CASE_1 = [
+    {
+        "data": [1, 2, 3, 4],
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[1, 3], [2, 4]],
+]
+
+TEST_CASE_2 = [
+    {
+        "data": [1, 2, 3, 4],
+        "num_partitions": 2,
+        "shuffle": True,
+        "seed": 123,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[4, 2], [1, 3]],
+]
+
+TEST_CASE_3 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[1, 3, 5], [2, 4]],
+]
+
+TEST_CASE_4 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": True,
+    },
+    [[1, 3, 5], [2, 4, 1]],
+]
+
+TEST_CASE_5 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": True,
+        "even_divisible": True,
+    },
+    [[1, 3], [2, 4]],
+]
+
+TEST_CASE_6 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "ratios": [3, 2],
+        "num_partitions": None,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": True,
+        "even_divisible": True,
+    },
+    [[1, 2, 3], [4, 5]],
+]
+
+TEST_CASE_7 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "ratios": [2, 1],
+        "num_partitions": None,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": True,
+        "even_divisible": True,
+    },
+    [[1, 2, 3], [4, 5]],
+]
+
+TEST_CASE_8 = [
+    {
+        "data": [1, 2, 3, 4, 5],
+        "ratios": [2, 1],
+        "num_partitions": None,
+        "shuffle": True,
+        "seed": 123,
+        "drop_last": True,
+        "even_divisible": True,
+    },
+    [[2, 4, 5], [1, 3]],
+]
+
+
+class TestPartitionDataset(unittest.TestCase):
+    @parameterized.expand(
+        [
+            TEST_CASE_1,
+            TEST_CASE_2,
+            TEST_CASE_3,
+            TEST_CASE_4,
+            TEST_CASE_5,
+            TEST_CASE_6,
+            TEST_CASE_7,
+            TEST_CASE_8,
+        ]
+    )
+    def test_value(self, input_param, result):
+        self.assertListEqual(partition_dataset(**input_param), result)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_partition_dataset_classes.py b/tests/test_partition_dataset_classes.py
new file mode 100644
--- /dev/null
+++ b/tests/test_partition_dataset_classes.py
@@ -0,0 +1,82 @@
+# 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
+
+from parameterized import parameterized
+
+from monai.data import partition_dataset_classes
+
+TEST_CASE_1 = [
+    {
+        "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+        "classes": [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3],
+        "ratios": [2, 1],
+        "num_partitions": None,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[2, 8, 4, 1, 3, 6, 5, 11, 12], [10, 13, 7, 9, 14]],
+]
+
+TEST_CASE_2 = [
+    {
+        "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+        "classes": [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3],
+        "ratios": None,
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[2, 10, 4, 1, 6, 9, 5, 12], [8, 13, 3, 7, 11, 14]],
+]
+
+TEST_CASE_3 = [
+    {
+        "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+        "classes": [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3],
+        "ratios": None,
+        "num_partitions": 2,
+        "shuffle": False,
+        "seed": 0,
+        "drop_last": False,
+        "even_divisible": True,
+    },
+    [[2, 10, 4, 1, 6, 9, 5, 12], [8, 2, 13, 3, 7, 1, 11, 14]],
+]
+
+TEST_CASE_4 = [
+    {
+        "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14],
+        "classes": [2, 0, 2, 1, 3, 2, 2, 0, 2, 0, 3, 3, 1, 3],
+        "ratios": [1, 2],
+        "num_partitions": None,
+        "shuffle": True,
+        "seed": 123,
+        "drop_last": False,
+        "even_divisible": False,
+    },
+    [[13, 7, 14, 2, 3], [11, 8, 6, 5, 9, 10, 12, 1, 4]],
+]
+
+
+class TestPartitionDatasetClasses(unittest.TestCase):
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
+    def test_value(self, input_param, result):
+        self.assertListEqual(partition_dataset_classes(**input_param), result)
+
+
+if __name__ == "__main__":
+    unittest.main()