[Python][ML] Adapt tutorials to the new RDataLoader interface

Author: siliataider

tutorials/machine_learning/ml_dataloader_NumPy.py

@@ -14,20 +14,16 @@
 file_name = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
 
 batch_size = 128
-chunk_size = 5000
-block_size = 400
 
 rdataframe = ROOT.RDataFrame(tree_name, file_name)
 
 target = "Type"
 
 num_of_epochs = 2
 
-gen_train, gen_validation = ROOT.Experimental.ML.CreateNumPyGenerators(
+gen_train, gen_validation = ROOT.Experimental.ML.RDataLoader(
     rdataframe,
     batch_size,
-    chunk_size,
-    block_size,
     target=target,
     validation_split=0.3,
     shuffle=True,
@@ -36,9 +32,9 @@
 
 for i in range(num_of_epochs):
     # Loop through training set
-    for i, (x_train, y_train) in enumerate(gen_train):
+    for i, (x_train, y_train) in enumerate(gen_train.AsNumpy()):
         print(f"Training batch {i + 1} => x: {x_train.shape}, y: {y_train.shape}")
 
     # Loop through Validation set
-    for i, (x_validation, y_validation) in enumerate(gen_validation):
+    for i, (x_validation, y_validation) in enumerate(gen_validation.AsNumpy()):
         print(f"Validation batch {i + 1} => x: {x_validation.shape}, y: {y_validation.shape}")

tutorials/machine_learning/ml_dataloader_PyTorch.py

@@ -15,20 +15,16 @@
 file_name = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
 
 batch_size = 128
-chunk_size = 5000
-block_size = 300
 
 rdataframe = ROOT.RDataFrame(tree_name, file_name)
 
 target = "Type"
 
 # Returns two generators that return training and validation batches
 # as PyTorch tensors.
-gen_train, gen_validation = ROOT.Experimental.ML.CreatePyTorchGenerators(
+gen_train, gen_validation = ROOT.Experimental.ML.RDataLoader(
     rdataframe,
     batch_size,
-    chunk_size,
-    block_size,
     target=target,
     validation_split=0.3,
     shuffle=True,
@@ -64,7 +60,7 @@ def calc_accuracy(targets, pred):
     print("Epoch ", i)
     model.train()
     # Loop through the training set and train model
-    for i, (x_train, y_train) in enumerate(gen_train):
+    for i, (x_train, y_train) in enumerate(gen_train.AsTorch()):
         # Make prediction and calculate loss
         pred = model(x_train)
         loss = loss_fn(pred, y_train)
@@ -85,7 +81,7 @@ def calc_accuracy(targets, pred):
 
     model.eval()
     # Evaluate the model on the validation set
-    for i, (x_val, y_val) in enumerate(gen_validation):
+    for i, (x_val, y_val) in enumerate(gen_validation.AsTorch()):
         # Make prediction and calculate accuracy
         pred = model(x_val)
         accuracy = calc_accuracy(y_val, pred)

tutorials/machine_learning/ml_dataloader_TensorFlow.py

@@ -18,18 +18,16 @@
 file_name = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
 
 batch_size = 128
-chunk_size = 5000
-block_size = 300
+approx_batches_in_memory = 50
 
 rdataframe = ROOT.RDataFrame(tree_name, file_name)
 target = ["Type"]
 
 # Returns two TF.Dataset for training and validation batches.
-ds_train, ds_valid = ROOT.Experimental.ML.CreateTFDatasets(
+ds_train, ds_valid = ROOT.Experimental.ML.RDataLoader(
     rdataframe,
     batch_size,
-    chunk_size,
-    block_size,
+    approx_batches_in_memory,
     target=target,
     validation_split=0.3,
     shuffle=True,
@@ -39,8 +37,8 @@
 num_of_epochs = 2
 
 # Datasets have to be repeated as many times as there are epochs
-ds_train_repeated = ds_train.repeat(num_of_epochs)
-ds_valid_repeated = ds_valid.repeat(num_of_epochs)
+ds_train_repeated = ds_train.AsTensorFlow().repeat(num_of_epochs)
+ds_valid_repeated = ds_valid.AsTensorFlow().repeat(num_of_epochs)
 
 # Number of batches per epoch must be given for model.fit
 train_batches_per_epoch = ds_train.number_of_batches

tutorials/machine_learning/ml_dataloader_filters_vectors.py

@@ -16,9 +16,7 @@
 tree_name = "test_tree"
 file_name = ROOT.gROOT.GetTutorialDir().Data() + "/machine_learning/ml_dataloader_filters_vectors_hvector.root"
 
-chunk_size = 50  # Defines the size of the chunks
 batch_size = 5  # Defines the size of the returned batches
-block_size = 10  # Defines the size of the blocks that builds up a chunk
 
 rdataframe = ROOT.RDataFrame(tree_name, file_name)
 
@@ -29,20 +27,18 @@
 
 max_vec_sizes = {"f4": 3, "f5": 2, "f6": 1}
 
-ds_train, ds_validation = ROOT.Experimental.ML.CreateNumPyGenerators(
+ds_train, ds_validation = ROOT.Experimental.ML.RDataLoader(
     filteredrdf,
     batch_size,
-    chunk_size,
-    block_size,
     validation_split=0.3,
     max_vec_sizes=max_vec_sizes,
     shuffle=False,
 )
 
 print(f"Columns: {ds_train.columns}")
 
-for i, b in enumerate(ds_train):
+for i, b in enumerate(ds_train.AsNumpy()):
     print(f"Training batch {i} => {b.shape}")
 
-for i, b in enumerate(ds_validation):
+for i, b in enumerate(ds_validation.AsNumpy()):
     print(f"Validation batch {i} => {b.shape}")