Skip to content

[Draft][ML] Data Loader: Async lazy loading #21299

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
siliataider wants to merge 3 commits into from
Closed

[Draft][ML] Data Loader: Async lazy loading #21299

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
dea7b2e
add py dataloader
siliataider Feb 13, 2026
ecf4723
add cpp batchgenerator
siliataider Feb 13, 2026
453b702
update ml_dataloader.py
siliataider Feb 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
# Author: Nopphakorn Subsa-Ard, King Mongkut's University of Technology Thonburi (KMUTT) (TH) 08/2024
# Author: Vincenzo Eduardo Padulano, CERN 10/2024
# Author: Martin Føll, University of Oslo (UiO) & CERN 01/2026
# Author: Silia Taider, CERN 02/2026

################################################################################
# Copyright (C) 1995-2026, Rene Brun and Fons Rademakers. #
Expand Down Expand Up @@ -501,6 +502,8 @@ def GetValidationBatch(self) -> Any:
class LoadingThreadContext:
def __init__(self, base_generator: BaseGenerator):
self.base_generator = base_generator
# init loading thread
self.base_generator.Activate()
# create training batches from the first chunk
self.base_generator.CreateTrainBatches()

Expand All @@ -509,7 +512,7 @@ def __enter__(self):

def __exit__(self, type, value, traceback):
self.base_generator.DeActivateTrainingEpoch()
return True
return False


class TrainDataLoader:
Expand Down Expand Up @@ -593,6 +596,8 @@ def __call__(self) -> Any:
class LoadingThreadContextVal:
def __init__(self, base_generator: BaseGenerator):
self.base_generator = base_generator
# init loading thread
self.base_generator.Activate()
# create validation batches from the first chunk
self.base_generator.CreateValidationBatches()

Expand All @@ -601,7 +606,7 @@ def __enter__(self):

def __exit__(self, type, value, traceback):
self.base_generator.DeActivateValidationEpoch()
return True
return False


class ValidationDataLoader:
Expand Down Expand Up @@ -668,7 +673,6 @@ def __call__(self) -> Any:
while True:
batch = self.base_generator.GetValidationBatch()
if batch is None:
self.base_generator.DeActivateValidationEpoch()
break
yield self.conversion_function(batch)

Expand Down
6 changes: 3 additions & 3 deletions bindings/pyroot/pythonizations/test/ml_dataloader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -540,7 +540,7 @@ def test09_filtered_last_chunk(self):

dff = df.Filter("b1 % 2 == 0", "name")

gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
dff,
batch_size=3,
chunk_size=9,
Expand Down Expand Up @@ -1622,7 +1622,7 @@ def test09_filtered_last_chunk(self):

dff = df.Filter("b1 % 2 == 0", "name")

gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
dff, batch_size=3, target="b2", validation_split=0, shuffle=False, drop_remainder=False, load_eager=True
)

Expand Down Expand Up @@ -2784,7 +2784,7 @@ def test09_filtered_last_chunk(self):
dff1 = df1.Filter("b1 % 2 == 0", "name")
dff2 = df2.Filter("b1 % 2 == 0", "name")

gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
[dff1, dff2],
batch_size=3,
target="b2",
Expand Down
Loading
Loading