diff --git a/monai/data/utils.py b/monai/data/utils.py index 5f1f8f9388..93da67e7b8 100644 --- a/monai/data/utils.py +++ b/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 diff --git a/tests/data/test_partition_dataset.py b/tests/data/test_partition_dataset.py index c93a6c7682..8259584005 100644 --- a/tests/data/test_partition_dataset.py +++ b/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()