diff --git a/monai/data/utils.py b/monai/data/utils.py index d548ed72486..d85b241d6da 100644 --- a/monai/data/utils.py +++ b/monai/data/utils.py @@ -1203,9 +1203,22 @@ def partition_dataset( if ratios: next_idx = 0 rsum = sum(ratios) - for r in ratios: + sizes = [int(r / rsum * data_len + 0.5) for r in ratios] + if sum(sizes) < data_len: + # Independent rounding can leave trailing items in no partition + # (e.g. ratios [1, 1, 1] on 10 items rounds to 3 + 3 + 3); assign + # the remainder so every item lands in exactly one partition. + # It goes to the last partition with a positive ratio: a + # zero-ratio partition must stay empty (e.g. ratios + # [1, 1, 1, 0] on 10 items rounds to 3 + 3 + 3 + 0, and the + # leftover item belongs in the third partition, not the empty one). + for i in range(len(ratios) - 1, -1, -1): + if ratios[i] > 0: + sizes[i] += data_len - sum(sizes) + break + 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 c93a6c76821..57e34e71b7b 100644 --- a/tests/data/test_partition_dataset.py +++ b/tests/data/test_partition_dataset.py @@ -117,10 +117,89 @@ ] -class TestPartitionDataset(unittest.TestCase): +TEST_CASE_9 = [ + { + "data": list(range(10)), + "ratios": [1, 1, 1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9]], +] + +TEST_CASE_10 = [ + { + "data": list(range(63)), + "ratios": [0.8, 0.1, 0.1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [list(range(50)), list(range(50, 56)), list(range(56, 63))], +] + +TEST_CASE_11 = [ + { + "data": list(range(32)), + "ratios": [1, 1, 1, 1, 1], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [list(range(6)), list(range(6, 12)), list(range(12, 18)), list(range(18, 24)), list(range(24, 32))], +] + +TEST_CASE_12 = [ + { + "data": [1, 2, 3, 4, 5], + "ratios": [0.6, 0.2, 0.2], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[1, 2, 3], [4], [5]], +] + +TEST_CASE_13 = [ + { + "data": list(range(10)), + "ratios": [1, 1, 1, 0], + "num_partitions": None, + "shuffle": False, + "seed": 0, + "drop_last": True, + "even_divisible": True, + }, + [[0, 1, 2], [3, 4, 5], [6, 7, 8, 9], []], +] + +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] + [ + TEST_CASE_1, + TEST_CASE_2, + TEST_CASE_3, + TEST_CASE_4, + TEST_CASE_5, + TEST_CASE_6, + TEST_CASE_7, + TEST_CASE_8, + TEST_CASE_9, + TEST_CASE_10, + TEST_CASE_11, + TEST_CASE_12, + TEST_CASE_13, + ] ) def test_value(self, input_param, result): self.assertListEqual(partition_dataset(**input_param), result)