partition_dataset #9220

When partition_dataset(data, ratios=...) splits a list, it computes every partition's size by rounding independently (see monai/data/utils.py, lines 1217–1224 at commit 6920610). If those rounded sizes add up to less than len(data), the extra items end up in no partition, and no warning or error is raised.

With three ratios this affects roughly one length in five: for 0.8/0.1/0.1 it happens for 200 of the lengths from 1 to 1000, and for 1/1/1 it happens for 334 of them. With five equal ratios, as many as two items can disappear. Because partition_dataset_classes applies the same function to each class, it can lose up to one item per class. Splits into two parts are not affected.

I found this by reading the code and then testing it.

from monai.data import partition_dataset

print(partition_dataset(list(range(10)), ratios=[1, 1, 1]))

data = list(range(63)) # e.g. the 63 training cases of MSD Task06_Lung
train, val, test = partition_dataset(data, ratios=[0.8, 0.1, 0.1], shuffle=True, seed=0)
print(len(train), len(val), len(test), "in no split:", sorted(set(data) - set(train + val + test)))

Output on dev (6920610):

[[0, 1, 2], [3, 4, 5], [6, 7, 8]]
50 6 6 in no split: [44]

The same thing occurs for other Decathlon training-set sizes. Examples: 303 cases (Task08_HepaticVessel) and 484 cases (Task01_BrainTumour) with 0.8/0.1/0.1; 63 and 303 cases with 0.7/0.15/0.15; and 32 cases (Task05_Prostate) with five equal ratios, where two cases are lost. In a Decathlon-style pipeline (load_decathlon_datalist → partition_dataset → CacheDataset → DataLoader) with 63 cases, one case is never loaded by the train, validation or test split.

Expected behavior

Each item should land in exactly one partition, as the docstring example shows ([0.6, 0.2, 0.2] on 5 items gives [[1, 2, 3], [4], [5]]). The documentation says drop_last does not apply when ratios are used.

Environment

MONAI 1.6.1rc0+50.g6920610c (dev), PyTorch 2.14.1+cpu, NumPy 2.5.4, Python 3.12.13, Linux.

Additional context

I have a small fix ready. It hands the leftover items to the partitions that were rounded down the most, and splits that already cover every item stay unchanged. It comes with tests, and I can open a PR.