fix(data): keep every item when partition_dataset rounds ratio sizes

monai/data/utils.py

@@ -1181,7 +1181,8 @@ def partition_dataset(
 
     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].
+        ratios: a list of ratio number to split the dataset, like [8, 1, 1]. If rounding the partition sizes
+            leaves some items over, they are added to the partitions that were rounded down the most.
         num_partitions: expected number of the partitions to evenly split, only works when `ratios` not specified.
         shuffle: whether to shuffle the original dataset before splitting.
         seed: random seed to shuffle the dataset, only works when `shuffle` is True.
@@ -1217,9 +1218,16 @@ def partition_dataset(
     if ratios:
         next_idx = 0
         rsum = sum(ratios)
-        for r in ratios:
+        exact = [r / rsum * data_len for r in ratios]
+        sizes = [int(e + 0.5) for e in exact]
+        # rounding each size on its own can leave items without a partition,
+        # give them to the partitions whose sizes were rounded down the most
+        missing = max(data_len - sum(sizes), 0)
+        for j in sorted(range(len(sizes)), key=lambda k: sizes[k] - exact[k])[:missing]:
+            sizes[j] += 1
+        for size in sizes:
             start_idx = next_idx
-            next_idx = min(start_idx + int(r / rsum * data_len + 0.5), data_len)
+            next_idx = min(start_idx + size, data_len)
             datasets.append([data[i] for i in indices[start_idx:next_idx]])
         return datasets
 

tests/data/test_partition_dataset.py

@@ -116,6 +116,12 @@
     [[2, 4, 5], [1, 3]],
 ]
 
+# rounding each partition size on its own leaves one item over in these cases
+TEST_CASES_ROUNDING = [
+    [{"data": list(range(1, 11)), "ratios": [1, 1, 1]}, [[1, 2, 3, 4], [5, 6, 7], [8, 9, 10]]],
+    [{"data": list(range(1, 14)), "ratios": [0.8, 0.1, 0.1]}, [list(range(1, 12)), [12], [13]]],
+]
+
 
 class TestPartitionDataset(unittest.TestCase):
 
@@ -125,6 +131,28 @@ class TestPartitionDataset(unittest.TestCase):
     def test_value(self, input_param, result):
         self.assertListEqual(partition_dataset(**input_param), result)
 
+    @parameterized.expand(TEST_CASES_ROUNDING)
+    def test_ratios_rounding(self, input_param, result):
+        """Items left over after rounding go to the partitions that were rounded down the most.
+
+        Args:
+            input_param: keyword arguments for ``partition_dataset``.
+            result: the expected partitions.
+        """
+        self.assertListEqual(partition_dataset(**input_param), result)
+
+    @parameterized.expand([(63, [0.8, 0.1, 0.1]), (303, [0.7, 0.15, 0.15]), (32, [1, 1, 1, 1, 1])])
+    def test_ratios_assign_every_item(self, data_len, ratios):
+        """Every item lands in exactly one partition when the data is shuffled.
+
+        Args:
+            data_len: number of items to partition.
+            ratios: the partition ratios.
+        """
+        data = list(range(data_len))
+        partitions = partition_dataset(data, ratios=ratios, shuffle=True, seed=0)
+        self.assertListEqual(sorted(i for part in partitions for i in part), data)
+
 
 if __name__ == "__main__":
     unittest.main()