From 721516d26d048c2c9e1a15a8151eacefd3380418 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Fri, 16 Oct 2026 14:55:08 +0200 Subject: [PATCH 1/8] [ML] Opt.: Shuffle the rows while creating the batches The batches read the rows of a chunk or dataset in shuffled order, instead of copying it into a shuffled tensor first. The leftover rows are appended to the leftover batch directly, which is moved to the queue once full. Assisted-by: Claude Opus 5.5 --- tree/ml/inc/ROOT/ML/RFlat2DMatrix.hxx | 2 + tree/ml/src/RBatchLoader.cxx | 91 +++++++------------------- tree/ml/src/RBatchLoader.hxx | 6 +- tree/ml/src/RDataLoaderEngine.cxx | 28 ++++---- tree/ml/src/RFlat2DMatrixOperators.cxx | 54 +++++++++++---- tree/ml/src/RFlat2DMatrixOperators.hxx | 8 +++ 6 files changed, 96 insertions(+), 93 deletions(-) diff --git a/tree/ml/inc/ROOT/ML/RFlat2DMatrix.hxx b/tree/ml/inc/ROOT/ML/RFlat2DMatrix.hxx index 9ebeac0e48b63..85dd2884ea7b0 100644 --- a/tree/ml/inc/ROOT/ML/RFlat2DMatrix.hxx +++ b/tree/ml/inc/ROOT/ML/RFlat2DMatrix.hxx @@ -46,6 +46,8 @@ struct RFlat2DMatrix { fRVec.resize(rows * cols); } + void Reserve(std::size_t rows, std::size_t cols) { fRVec.reserve(rows * cols); } + void Reshape(std::size_t rows, std::size_t cols) { // We don't reallocate: require matching sizes diff --git a/tree/ml/src/RBatchLoader.cxx b/tree/ml/src/RBatchLoader.cxx index a8fa10ed693fa..a6caef30c34c7 100644 --- a/tree/ml/src/RBatchLoader.cxx +++ b/tree/ml/src/RBatchLoader.cxx @@ -5,6 +5,8 @@ #include #include +#include "RFlat2DMatrixOperators.hxx" + namespace ROOT::Experimental::Internal::ML { RBatchLoader::RBatchLoader(std::size_t batchSize, const std::vector &cols, std::mutex &sharedMutex, @@ -24,7 +26,6 @@ RBatchLoader::RBatchLoader(std::size_t batchSize, const std::vector RecalculateBatchCounts(numEntries); fPrimaryLeftoverBatch = std::make_unique(); - fSecondaryLeftoverBatch = std::make_unique(); } /// \brief Activate the batchloader. This means that batches can be created and loaded. @@ -67,7 +68,6 @@ void RBatchLoader::Reset() fCurrentBatch.reset(); fPrimaryLeftoverBatch = std::make_unique(); - fSecondaryLeftoverBatch = std::make_unique(); } fCV.notify_all(); @@ -83,13 +83,15 @@ void RBatchLoader::MarkProducerDone() /// \brief Return a batch of data as a unique pointer. /// After the batch has been processed, it should be destroyed. /// \param[in] chunkTensor Tensor with the data from the chunk +/// \param[in] permutation Order in which to read the rows of the chunk, empty to read them in order /// \param[in] idxs Index of batch in the chunk /// \return Batch -std::unique_ptr RBatchLoader::CreateBatch(RFlat2DMatrix &chunkTensor, std::size_t idxs) +std::unique_ptr RBatchLoader::CreateBatch(const RFlat2DMatrix &chunkTensor, + const std::vector &permutation, std::size_t idxs) { - auto batch = std::make_unique(fBatchSize, fNumColumns); - std::copy(chunkTensor.GetData() + (idxs * fBatchSize * fNumColumns), - chunkTensor.GetData() + ((idxs + 1) * fBatchSize * fNumColumns), batch->GetData()); + auto batch = std::make_unique(); + batch->Reserve(fBatchSize, fNumColumns); + RFlat2DMatrixOperators::AppendRows(*batch, chunkTensor, permutation, idxs * fBatchSize, (idxs + 1) * fBatchSize); return batch; } @@ -123,86 +125,43 @@ RFlat2DMatrix RBatchLoader::GetBatch() /// \brief Creating the batches from a chunk and add them to the queue. /// \param[in] chunkTensor Tensor with the data from the chunk +/// \param[in] permutation Order in which to read the rows of the chunk, empty to read them in order /// \param[in] isLastBatch Check if the batch in the chunk is the last one -void RBatchLoader::CreateBatches(RFlat2DMatrix &chunkTensor, bool isLastBatch) +void RBatchLoader::CreateBatches(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, + bool isLastBatch) { std::size_t chunkSize = chunkTensor.GetRows(); - std::size_t numCols = chunkTensor.GetCols(); std::size_t numBatches = chunkSize / fBatchSize; - std::size_t leftoverBatchSize = chunkSize % fBatchSize; // create a vector of batches std::vector> batches; // fill the full batches from the chunk into a vector for (std::size_t i = 0; i < numBatches; i++) { - batches.emplace_back(CreateBatch(chunkTensor, i)); + batches.emplace_back(CreateBatch(chunkTensor, permutation, i)); } - // copy the remaining entries from the chunk into a leftover batch - RFlat2DMatrix LeftoverBatch(leftoverBatchSize, numCols); - std::copy(chunkTensor.GetData() + (numBatches * fBatchSize * numCols), - chunkTensor.GetData() + (numBatches * fBatchSize * numCols + leftoverBatchSize * numCols), - LeftoverBatch.GetData()); - - // calculate how many empty slots are left in fPrimaryLeftoverBatch - std::size_t PrimaryLeftoverSize = fPrimaryLeftoverBatch->GetRows(); - std::size_t emptySlots = fBatchSize - PrimaryLeftoverSize; - - // copy LeftoverBatch to end of fPrimaryLeftoverBatch - if (emptySlots >= leftoverBatchSize) { - fPrimaryLeftoverBatch->Resize(PrimaryLeftoverSize + leftoverBatchSize, numCols); - std::copy(LeftoverBatch.GetData(), LeftoverBatch.GetData() + (leftoverBatchSize * fNumColumns), - fPrimaryLeftoverBatch->GetData() + (PrimaryLeftoverSize * numCols)); - - // copy LeftoverBatch to end of fPrimaryLeftoverBatch and add it to the batch - if (emptySlots == leftoverBatchSize) { - auto copy = std::make_unique(fBatchSize, fNumColumns); - std::copy(fPrimaryLeftoverBatch->GetData(), fPrimaryLeftoverBatch->GetData() + (fBatchSize * fNumColumns), - copy->GetData()); - batches.emplace_back(std::move(copy)); - - // reset fPrimaryLeftoverBatch and fSecondaryLeftoverBatch - *fPrimaryLeftoverBatch = *fSecondaryLeftoverBatch; - fSecondaryLeftoverBatch = std::make_unique(); - } - } + // append the remaining entries of the chunk to fPrimaryLeftoverBatch, and add it to the batches once it is full + for (std::size_t row = numBatches * fBatchSize; row < chunkSize;) { + const std::size_t numRows = std::min(chunkSize - row, fBatchSize - fPrimaryLeftoverBatch->GetRows()); + fPrimaryLeftoverBatch->Reserve(fBatchSize, fNumColumns); + RFlat2DMatrixOperators::AppendRows(*fPrimaryLeftoverBatch, chunkTensor, permutation, row, row + numRows); + row += numRows; - // copy LeftoverBatch to both fPrimaryLeftoverBatch and fSecondaryLeftoverBatch - else if (emptySlots < leftoverBatchSize) { - // copy the first part of LeftoverBatch to end of fPrimaryLeftoverTrainingBatch - fPrimaryLeftoverBatch->Resize(fBatchSize, numCols); - std::copy(LeftoverBatch.GetData(), LeftoverBatch.GetData() + (emptySlots * numCols), - fPrimaryLeftoverBatch->GetData() + (PrimaryLeftoverSize * numCols)); - - // copy the last part of LeftoverBatch to the end of fSecondaryLeftoverBatch - fSecondaryLeftoverBatch->Resize(leftoverBatchSize - emptySlots, numCols); - std::copy(LeftoverBatch.GetData() + (emptySlots * numCols), - LeftoverBatch.GetData() + (leftoverBatchSize * numCols), fSecondaryLeftoverBatch->GetData()); - - // add fPrimaryLeftoverBatch to the batch vector - auto copy = std::make_unique(fBatchSize, fNumColumns); - std::copy(fPrimaryLeftoverBatch->GetData(), fPrimaryLeftoverBatch->GetData() + (fBatchSize * fNumColumns), - copy->GetData()); - batches.emplace_back(std::move(copy)); - - // exchange fPrimaryLeftoverBatch and fSecondaryLeftoverBatch - *fPrimaryLeftoverBatch = *fSecondaryLeftoverBatch; - // reset fSecondaryLeftoverTrainingBatch - fSecondaryLeftoverBatch = std::make_unique(); + if (fPrimaryLeftoverBatch->GetRows() == fBatchSize) { + batches.emplace_back(std::move(fPrimaryLeftoverBatch)); + fPrimaryLeftoverBatch = std::make_unique(); + } } - // copy the content of fPrimaryLeftoverBatch to the leftover batch from the chunk + // add the content of fPrimaryLeftoverBatch as the leftover batch of the last chunk if (isLastBatch) { if (!fDropRemainder && fLeftoverBatchSize > 0) { - auto copy = std::make_unique(fLeftoverBatchSize, fNumColumns); - std::copy(fPrimaryLeftoverBatch->GetData(), - fPrimaryLeftoverBatch->GetData() + (fLeftoverBatchSize * fNumColumns), copy->GetData()); - batches.emplace_back(std::move(copy)); + fPrimaryLeftoverBatch->Resize(fLeftoverBatchSize, fNumColumns); + batches.emplace_back(std::move(fPrimaryLeftoverBatch)); } fPrimaryLeftoverBatch = std::make_unique(); - fSecondaryLeftoverBatch = std::make_unique(); } { diff --git a/tree/ml/src/RBatchLoader.hxx b/tree/ml/src/RBatchLoader.hxx index 44575c42b9eea..70ad48cdd7532 100644 --- a/tree/ml/src/RBatchLoader.hxx +++ b/tree/ml/src/RBatchLoader.hxx @@ -63,7 +63,6 @@ private: // primary and secondary leftover batches used to create batches from a chunk std::unique_ptr fPrimaryLeftoverBatch; - std::unique_ptr fSecondaryLeftoverBatch; public: RBatchLoader(std::size_t batchSize, const std::vector &cols, std::mutex &sharedMutex, @@ -75,9 +74,10 @@ public: void Reset(); void MarkProducerDone(); - std::unique_ptr CreateBatch(RFlat2DMatrix &chunkTensor, std::size_t idxs); + std::unique_ptr + CreateBatch(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, std::size_t idxs); RFlat2DMatrix GetBatch(); - void CreateBatches(RFlat2DMatrix &chunkTensor, bool isLastBatch); + void CreateBatches(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, bool isLastBatch); void RecalculateBatchCounts(std::size_t numEntries); bool isProducerDone() { return fProducerDone; } diff --git a/tree/ml/src/RDataLoaderEngine.cxx b/tree/ml/src/RDataLoaderEngine.cxx index 9a2e75591f24e..7342903199574 100644 --- a/tree/ml/src/RDataLoaderEngine.cxx +++ b/tree/ml/src/RDataLoaderEngine.cxx @@ -334,7 +334,7 @@ void RDataLoaderEngine::LoadData() lock.unlock(); // reserve the buffer and let the clusters be appended to it, filters can only make them smaller RFlat2DMatrix stagingBuffer; - stagingBuffer.fRVec.reserve(accumulatedEntries * fClusterLoader->GetNumChunkCols()); + stagingBuffer.Reserve(accumulatedEntries, fClusterLoader->GetNumChunkCols()); std::size_t rowOffset = 0; for (auto &cluster : trainClustersToLoad) { @@ -357,9 +357,9 @@ void RDataLoaderEngine::LoadData() stagingBuffer.Resize(rowOffset, fClusterLoader->GetNumChunkCols()); - RFlat2DMatrix shuffledStagingBuffer; - fTrainingBatchLoader->CreateBatches(fTensorOperators->ShuffleTensor(shuffledStagingBuffer, stagingBuffer), - isLastBuffer); + // the batches read the buffer in shuffled order, so it needs no shuffled copy + fTrainingBatchLoader->CreateBatches( + stagingBuffer, fTensorOperators->MakePermutation(stagingBuffer.GetRows()), isLastBuffer); // Re-acquire the lock before the next iteration to check conditions and update indices lock.lock(); @@ -408,7 +408,7 @@ void RDataLoaderEngine::LoadData() lock.unlock(); RFlat2DMatrix stagingBuffer; - stagingBuffer.fRVec.reserve(accumulatedEntries * fClusterLoader->GetNumChunkCols()); + stagingBuffer.Reserve(accumulatedEntries, fClusterLoader->GetNumChunkCols()); std::size_t rowOffset = 0; for (const auto &cluster : valClustersToLoad) { @@ -417,9 +417,8 @@ void RDataLoaderEngine::LoadData() } stagingBuffer.Resize(rowOffset, fClusterLoader->GetNumChunkCols()); - RFlat2DMatrix shuffledStagingBuffer; - fValidationBatchLoader->CreateBatches(fTensorOperators->ShuffleTensor(shuffledStagingBuffer, stagingBuffer), - isLastBuffer); + fValidationBatchLoader->CreateBatches( + stagingBuffer, fTensorOperators->MakePermutation(stagingBuffer.GetRows()), isLastBuffer); lock.lock(); } @@ -435,15 +434,18 @@ void RDataLoaderEngine::CreateTrainBatches() if (fLoadEager) { RFlat2DMatrix *source = &fSampledTrainingDataset; + std::vector permutation; if (fSampleType == "") { - source = &fTensorOperators->ShuffleTensor(fSampledTrainingDataset, fTrainingDataset); + // the batches read the dataset in shuffled order, so it needs no shuffled copy + source = &fTrainingDataset; + permutation = fTensorOperators->MakePermutation(fTrainingDataset.GetRows()); } else { fTrainingSampler->Sampler(fSampledTrainingDataset); } - fTrainingBatchLoader->CreateBatches(*source, true); + fTrainingBatchLoader->CreateBatches(*source, permutation, true); fTrainingBatchLoader->MarkProducerDone(); } } @@ -456,15 +458,17 @@ void RDataLoaderEngine::CreateValidationBatches() if (fLoadEager) { RFlat2DMatrix *source = &fSampledValidationDataset; + std::vector permutation; if (fSampleType == "") { - source = &fTensorOperators->ShuffleTensor(fSampledValidationDataset, fValidationDataset); + source = &fValidationDataset; + permutation = fTensorOperators->MakePermutation(fValidationDataset.GetRows()); } else { fValidationSampler->Sampler(fSampledValidationDataset); } - fValidationBatchLoader->CreateBatches(*source, true); + fValidationBatchLoader->CreateBatches(*source, permutation, true); fValidationBatchLoader->MarkProducerDone(); } } diff --git a/tree/ml/src/RFlat2DMatrixOperators.cxx b/tree/ml/src/RFlat2DMatrixOperators.cxx index 20bcb077375a0..5e111ea076c69 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.cxx +++ b/tree/ml/src/RFlat2DMatrixOperators.cxx @@ -5,7 +5,6 @@ #include #include "ROOT/ML/RFlat2DMatrix.hxx" -#include "Rtypes.h" namespace ROOT::Experimental::Internal::ML { @@ -13,11 +12,25 @@ RFlat2DMatrixOperators::~RFlat2DMatrixOperators() = default; RFlat2DMatrix &RFlat2DMatrixOperators::ShuffleTensor(RFlat2DMatrix &ShuffledTensor, RFlat2DMatrix &Tensor) { + const auto permutation = MakePermutation(Tensor.GetRows()); // nothing to permute, no copy - if (!fShuffle) { + if (permutation.empty()) { return Tensor; } + ShuffledTensor.fRVec.clear(); + ShuffledTensor.fRVec.reserve(Tensor.GetSize()); + AppendRows(ShuffledTensor, Tensor, permutation, 0, Tensor.GetRows()); + + return ShuffledTensor; +} + +std::vector RFlat2DMatrixOperators::MakePermutation(std::size_t rows) +{ + if (!fShuffle) { + return {}; + } + std::random_device rd; std::mt19937 g; @@ -27,24 +40,41 @@ RFlat2DMatrix &RFlat2DMatrixOperators::ShuffleTensor(RFlat2DMatrix &ShuffledTens g.seed(fSetSeed); } - std::size_t rows = Tensor.GetRows(); - std::size_t cols = Tensor.GetCols(); - ShuffledTensor.Resize(rows, cols); - // make an identity permutation map - std::vector indices(rows); + std::vector indices(rows); std::iota(indices.begin(), indices.end(), 0); // shuffle the identity permutation to create a new permutation std::shuffle(indices.begin(), indices.end(), g); - // shuffle data in the tensor with the permutation map defined above - for (std::size_t i = 0; i < rows; i++) { - std::copy(Tensor.GetData() + indices[i] * cols, Tensor.GetData() + (indices[i] + 1) * cols, - ShuffledTensor.GetData() + i * cols); + return indices; +} + +void RFlat2DMatrixOperators::AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix &Tensor, + const std::vector &permutation, std::size_t begin, std::size_t end) +{ + if (begin == end) { + return; } - return ShuffledTensor; + const std::size_t cols = Tensor.GetCols(); + const float *data = Tensor.GetData(); + auto &values = Dest.fRVec; + + // grow once and copy the rows in place, cheaper than inserting them one by one + const std::size_t offset = values.size(); + values.resize(offset + (end - begin) * cols); + float *out = values.data() + offset; + + if (permutation.empty()) { + std::copy(data + begin * cols, data + end * cols, out); + } else { + for (std::size_t i = begin; i < end; i++) { + out = std::copy_n(data + permutation[i] * cols, cols, out); + } + } + + Dest.Reshape(values.size() / cols, cols); } void RFlat2DMatrixOperators::SliceTensor(RFlat2DMatrix &SlicedTensor, RFlat2DMatrix &Tensor, diff --git a/tree/ml/src/RFlat2DMatrixOperators.hxx b/tree/ml/src/RFlat2DMatrixOperators.hxx index 3c9f870c125e9..854afb90f593c 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.hxx +++ b/tree/ml/src/RFlat2DMatrixOperators.hxx @@ -11,6 +11,7 @@ #ifndef ROOT_INTERNAL_ML_RFLAT2DMATRIXOPERATORS #define ROOT_INTERNAL_ML_RFLAT2DMATRIXOPERATORS +#include #include // Forward decl @@ -38,6 +39,13 @@ public: // returns the tensor to read from: the permuted copy, or Tensor itself when shuffling is off RFlat2DMatrix &ShuffleTensor(RFlat2DMatrix &ShuffledTensor, RFlat2DMatrix &Tensor); + // returns the order in which to read the rows of a tensor, empty when shuffling is off + std::vector MakePermutation(std::size_t rows); + + // appends the rows [begin, end) of Tensor to Dest, read in the order of the permutation unless it is empty + static void AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix &Tensor, const std::vector &permutation, + std::size_t begin, std::size_t end); + void SliceTensor(RFlat2DMatrix &SlicedTensor, RFlat2DMatrix &Tensor, const std::vector> &slice); From 128d8ceb9e971a70ee9e9578892f17ef041bdd96 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Fri, 16 Oct 2026 14:56:47 +0200 Subject: [PATCH 2/8] [ML] Opt.: Gather the rows when splitting and sampling the eager datasets The training and validation datasets, and the sampled ones, are filled by reading the loaded rows in shuffled order, instead of building shuffled, sliced and concatenated copies first. The datasets of several dataframes are split directly into the concatenated ones. Assisted-by: Claude Opus 5.5 --- tree/ml/src/RDataLoaderEngine.cxx | 4 +- tree/ml/src/RDatasetLoader.cxx | 110 ++++++++++++++++--------- tree/ml/src/RDatasetLoader.hxx | 6 +- tree/ml/src/RFlat2DMatrixOperators.cxx | 52 ------------ tree/ml/src/RFlat2DMatrixOperators.hxx | 8 -- tree/ml/src/RSampler.cxx | 57 ++++++------- tree/ml/src/RSampler.hxx | 3 + 7 files changed, 101 insertions(+), 139 deletions(-) diff --git a/tree/ml/src/RDataLoaderEngine.cxx b/tree/ml/src/RDataLoaderEngine.cxx index 7342903199574..37b61fc13287d 100644 --- a/tree/ml/src/RDataLoaderEngine.cxx +++ b/tree/ml/src/RDataLoaderEngine.cxx @@ -77,11 +77,9 @@ RDataLoaderEngine::RDataLoaderEngine(const std::vector &rdfs, if (fLoadEager) { fDatasetLoader = std::make_unique(fRdfs, fTestSize, fCols, fVecSizes, vecPadding, fShuffle, fSetSeed); - fDatasetLoader->SplitDatasets(); + fDatasetLoader->SplitDatasets(fSampleType == ""); if (fSampleType == "") { - fDatasetLoader->ConcatenateDatasets(); - fTrainingDataset = fDatasetLoader->ReleaseTrainingDataset(); fValidationDataset = fDatasetLoader->ReleaseValidationDataset(); diff --git a/tree/ml/src/RDatasetLoader.cxx b/tree/ml/src/RDatasetLoader.cxx index b0b92c32cbd77..4d280be9bb55f 100644 --- a/tree/ml/src/RDatasetLoader.cxx +++ b/tree/ml/src/RDatasetLoader.cxx @@ -21,12 +21,10 @@ RDatasetLoader::RDatasetLoader(const std::vector &rdfs, const } ////////////////////////////////////////////////////////////////////////// -/// \brief Split an individual dataframe into a training and validation dataset -/// \param[in] rdf Dataframe that will be split into training and validation -/// \param[in] TrainingDataset Tensor for the training dataset -/// \param[in] ValidationDataset Tensor for the validation dataset -void RDatasetLoader::SplitDataframe(ROOT::RDF::RNode &rdf, RFlat2DMatrix &TrainingDataset, - RFlat2DMatrix &ValidationDataset) +/// \brief Load all the entries of an individual dataframe +/// \param[in] rdf Dataframe to load +/// \return Tensor with the dataset +RFlat2DMatrix RDatasetLoader::LoadDataframe(ROOT::RDF::RNode &rdf) { const bool NotFiltered = rdf.GetFilterNames().empty(); @@ -53,10 +51,6 @@ void RDatasetLoader::SplitDataframe(ROOT::RDF::RNode &rdf, RFlat2DMatrix &Traini Entries->push_back((*Entries)[NumEntries - 1] + 1); } - // number of training and validation entries after the split - std::size_t NumValidationEntries = static_cast(fValidationSplit * NumEntries); - std::size_t NumTrainingEntries = NumEntries - NumValidationEntries; - // reserve the dataset and let the values be appended to it RFlat2DMatrix Dataset; Dataset.fRVec.reserve(NumEntries * fNumDatasetCols); @@ -78,48 +72,86 @@ void RDatasetLoader::SplitDataframe(ROOT::RDF::RNode &rdf, RFlat2DMatrix &Traini ROOT::Internal::RDF::ChangeBeginAndEndEntries(rdf, (*Entries)[0], (*Entries)[NumEntries]); } - // copy out the validation tail, then shrink the (shuffled) buffer to the training rows and move it - RFlat2DMatrix ShuffledDataset; - RFlat2DMatrix &Source = fTensorOperators->ShuffleTensor(ShuffledDataset, Dataset); - fTensorOperators->SliceTensor(ValidationDataset, Source, {{NumTrainingEntries, NumEntries}, {0, fNumDatasetCols}}); - Source.Resize(NumTrainingEntries, fNumDatasetCols); - TrainingDataset = std::move(Source); + return Dataset; } ////////////////////////////////////////////////////////////////////////// -/// \brief Split the dataframes in a training and validation dataset -void RDatasetLoader::SplitDatasets() +/// \brief Split a loaded dataset into training and validation entries, appended to the given datasets +/// \param[in] Dataset Tensor with the loaded dataset +/// \param[in] TrainingDataset Tensor for the training dataset +/// \param[in] ValidationDataset Tensor for the validation dataset +void RDatasetLoader::SplitDataset(RFlat2DMatrix &Dataset, RFlat2DMatrix &TrainingDataset, + RFlat2DMatrix &ValidationDataset) { - fNumTrainingEntries = 0; - fNumValidationEntries = 0; - - for (auto &rdf : f_rdfs) { - RFlat2DMatrix TrainingDataset; - RFlat2DMatrix ValidationDataset; + // number of training and validation entries after the split + std::size_t NumEntries = Dataset.GetRows(); + std::size_t NumValidationEntries = static_cast(fValidationSplit * NumEntries); + std::size_t NumTrainingEntries = NumEntries - NumValidationEntries; - SplitDataframe(rdf, TrainingDataset, ValidationDataset); + // read the entries in shuffled order, instead of copying the dataset into a shuffled one first + const auto permutation = fTensorOperators->MakePermutation(NumEntries); - fNumTrainingEntries += TrainingDataset.GetRows(); - fNumValidationEntries += ValidationDataset.GetRows(); + ValidationDataset.fRVec.reserve(ValidationDataset.GetSize() + NumValidationEntries * fNumDatasetCols); + RFlat2DMatrixOperators::AppendRows(ValidationDataset, Dataset, permutation, NumTrainingEntries, NumEntries); - fTrainingDatasets.push_back(std::move(TrainingDataset)); - fValidationDatasets.push_back(std::move(ValidationDataset)); - } + TrainingDataset.fRVec.reserve(TrainingDataset.GetSize() + NumTrainingEntries * fNumDatasetCols); + RFlat2DMatrixOperators::AppendRows(TrainingDataset, Dataset, permutation, 0, NumTrainingEntries); } ////////////////////////////////////////////////////////////////////////// -/// \brief Concatenate the datasets to a dataset -void RDatasetLoader::ConcatenateDatasets() +/// \brief Split the dataframes in a training and validation dataset +/// \param[in] concatenate Whether to concatenate the datasets of all the dataframes into one +void RDatasetLoader::SplitDatasets(bool concatenate) { - if (fTrainingDatasets.size() == 1) { - fTrainingDataset = std::move(fTrainingDatasets[0]); - fValidationDataset = std::move(fValidationDatasets[0]); + std::vector Datasets; + for (auto &rdf : f_rdfs) { + Datasets.push_back(LoadDataframe(rdf)); + } + + if (concatenate && Datasets.size() > 1) { + // split every dataset directly into the concatenated ones, sized for all of them + std::size_t NumValidationEntries = 0; + std::size_t NumEntries = 0; + for (const auto &Dataset : Datasets) { + NumEntries += Dataset.GetRows(); + NumValidationEntries += static_cast(fValidationSplit * Dataset.GetRows()); + } + fTrainingDataset.fRVec.reserve((NumEntries - NumValidationEntries) * fNumDatasetCols); + fValidationDataset.fRVec.reserve(NumValidationEntries * fNumDatasetCols); + + for (auto &Dataset : Datasets) { + SplitDataset(Dataset, fTrainingDataset, fValidationDataset); + Dataset = RFlat2DMatrix(); + } + } else { + for (auto &Dataset : Datasets) { + RFlat2DMatrix TrainingDataset; + RFlat2DMatrix ValidationDataset; + SplitDataset(Dataset, TrainingDataset, ValidationDataset); + Dataset = RFlat2DMatrix(); + fTrainingDatasets.push_back(std::move(TrainingDataset)); + fValidationDatasets.push_back(std::move(ValidationDataset)); + } + + if (concatenate) { + fTrainingDataset = std::move(fTrainingDatasets[0]); + fValidationDataset = std::move(fValidationDatasets[0]); + fTrainingDatasets.clear(); + fValidationDatasets.clear(); + } + } + + fNumTrainingEntries = 0; + fNumValidationEntries = 0; + if (concatenate) { + fNumTrainingEntries = fTrainingDataset.GetRows(); + fNumValidationEntries = fValidationDataset.GetRows(); } else { - fTensorOperators->ConcatenateTensors(fTrainingDataset, fTrainingDatasets); - fTensorOperators->ConcatenateTensors(fValidationDataset, fValidationDatasets); + for (std::size_t i = 0; i < fTrainingDatasets.size(); i++) { + fNumTrainingEntries += fTrainingDatasets[i].GetRows(); + fNumValidationEntries += fValidationDatasets[i].GetRows(); + } } - fTrainingDatasets.clear(); - fValidationDatasets.clear(); } } // namespace ROOT::Experimental::Internal::ML diff --git a/tree/ml/src/RDatasetLoader.hxx b/tree/ml/src/RDatasetLoader.hxx index c0ef5d1717205..cb04093ba278d 100644 --- a/tree/ml/src/RDatasetLoader.hxx +++ b/tree/ml/src/RDatasetLoader.hxx @@ -65,11 +65,11 @@ public: const std::vector &cols, const std::vector &vecSizes = {}, const float vecPadding = 0.0, bool shuffle = true, const std::size_t setSeed = 0); - void SplitDataframe(ROOT::RDF::RNode &rdf, RFlat2DMatrix &TrainingDataset, RFlat2DMatrix &ValidationDataset); + RFlat2DMatrix LoadDataframe(ROOT::RDF::RNode &rdf); - void SplitDatasets(); + void SplitDataset(RFlat2DMatrix &Dataset, RFlat2DMatrix &TrainingDataset, RFlat2DMatrix &ValidationDataset); - void ConcatenateDatasets(); + void SplitDatasets(bool concatenate); std::vector ReleaseTrainingDatasets() { return std::move(fTrainingDatasets); } std::vector ReleaseValidationDatasets() { return std::move(fValidationDatasets); } diff --git a/tree/ml/src/RFlat2DMatrixOperators.cxx b/tree/ml/src/RFlat2DMatrixOperators.cxx index 5e111ea076c69..839e0b0cb2684 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.cxx +++ b/tree/ml/src/RFlat2DMatrixOperators.cxx @@ -10,21 +10,6 @@ namespace ROOT::Experimental::Internal::ML { RFlat2DMatrixOperators::~RFlat2DMatrixOperators() = default; -RFlat2DMatrix &RFlat2DMatrixOperators::ShuffleTensor(RFlat2DMatrix &ShuffledTensor, RFlat2DMatrix &Tensor) -{ - const auto permutation = MakePermutation(Tensor.GetRows()); - // nothing to permute, no copy - if (permutation.empty()) { - return Tensor; - } - - ShuffledTensor.fRVec.clear(); - ShuffledTensor.fRVec.reserve(Tensor.GetSize()); - AppendRows(ShuffledTensor, Tensor, permutation, 0, Tensor.GetRows()); - - return ShuffledTensor; -} - std::vector RFlat2DMatrixOperators::MakePermutation(std::size_t rows) { if (!fShuffle) { @@ -77,41 +62,4 @@ void RFlat2DMatrixOperators::AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix Dest.Reshape(values.size() / cols, cols); } -void RFlat2DMatrixOperators::SliceTensor(RFlat2DMatrix &SlicedTensor, RFlat2DMatrix &Tensor, - const std::vector> &slice) -{ - const auto &rowSlice = slice[0]; - const auto &colSlice = slice[1]; - - std::size_t rowStart = rowSlice[0]; - std::size_t rowEnd = rowSlice[1]; - std::size_t colStart = colSlice[0]; - std::size_t colEnd = colSlice[1]; - - std::size_t rows = rowEnd - rowStart; - std::size_t cols = colEnd - colStart; - - SlicedTensor.Resize(rows, cols); - std::copy(Tensor.GetData() + rowStart * cols, Tensor.GetData() + rowStart * cols + rows * cols, - SlicedTensor.GetData()); -} - -void RFlat2DMatrixOperators::ConcatenateTensors(RFlat2DMatrix &ConcatTensor, const std::vector &Tensors) -{ - std::size_t cols = Tensors[0].GetCols(); - std::size_t rows = 0; - - for (const auto &t : Tensors) { - rows += t.GetRows(); - } - - ConcatTensor.Resize(rows, cols); - - std::size_t index = 0; - for (std::size_t i = 0; i < Tensors.size(); i++) { - std::size_t tensorRows = Tensors[i].GetRows(); - std::copy(Tensors[i].GetData(), Tensors[i].GetData() + tensorRows * cols, ConcatTensor.GetData() + index * cols); - index += tensorRows; - } -} } // namespace ROOT::Experimental::Internal::ML diff --git a/tree/ml/src/RFlat2DMatrixOperators.hxx b/tree/ml/src/RFlat2DMatrixOperators.hxx index 854afb90f593c..12dbd3f65a6ba 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.hxx +++ b/tree/ml/src/RFlat2DMatrixOperators.hxx @@ -36,20 +36,12 @@ public: ~RFlat2DMatrixOperators(); - // returns the tensor to read from: the permuted copy, or Tensor itself when shuffling is off - RFlat2DMatrix &ShuffleTensor(RFlat2DMatrix &ShuffledTensor, RFlat2DMatrix &Tensor); - // returns the order in which to read the rows of a tensor, empty when shuffling is off std::vector MakePermutation(std::size_t rows); // appends the rows [begin, end) of Tensor to Dest, read in the order of the permutation unless it is empty static void AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix &Tensor, const std::vector &permutation, std::size_t begin, std::size_t end); - - void - SliceTensor(RFlat2DMatrix &SlicedTensor, RFlat2DMatrix &Tensor, const std::vector> &slice); - - void ConcatenateTensors(RFlat2DMatrix &ConcatTensor, const std::vector &Tensors); }; } // namespace ROOT::Experimental::Internal::ML diff --git a/tree/ml/src/RSampler.cxx b/tree/ml/src/RSampler.cxx index dd8894ee4a31c..6285d48b13996 100644 --- a/tree/ml/src/RSampler.cxx +++ b/tree/ml/src/RSampler.cxx @@ -107,24 +107,7 @@ void RSampler::RandomUndersampler(RFlat2DMatrix &ShuffledTensor) SampleWithoutReplacement(fNumResampledMajor, fNumMajor); } - std::size_t cols = fDatasets[0].GetCols(); - RFlat2DMatrix SampledTensor(fNumEntries, cols); - RFlat2DMatrix UndersampledMajorTensor(fNumResampledMajor, cols); - - std::size_t index = 0; - for (std::size_t i = 0; i < fNumResampledMajor; i++) { - std::copy(fDatasets[fMajor].GetData() + fSamples[i] * cols, - fDatasets[fMajor].GetData() + (fSamples[i] + 1) * cols, - UndersampledMajorTensor.GetData() + index * cols); - index++; - } - - fTensorOperators->ConcatenateTensors(SampledTensor, {UndersampledMajorTensor, fDatasets[fMinor]}); - if (fShuffle) { - fTensorOperators->ShuffleTensor(ShuffledTensor, SampledTensor); - } else { - ShuffledTensor = std::move(SampledTensor); - } + GatherSamples(ShuffledTensor, fDatasets[fMajor], fNumResampledMajor, fDatasets[fMinor]); } ////////////////////////////////////////////////////////////////////////// @@ -134,23 +117,29 @@ void RSampler::RandomOversampler(RFlat2DMatrix &ShuffledTensor) { SampleWithReplacement(fNumResampledMinor, fNumMinor); - std::size_t cols = fDatasets[0].GetCols(); - RFlat2DMatrix SampledTensor(fNumEntries, cols); - RFlat2DMatrix OversampledMinorTensor(fNumResampledMinor, cols); - - std::size_t index = 0; - for (std::size_t i = 0; i < fNumResampledMinor; i++) { - std::copy(fDatasets[fMinor].GetData() + fSamples[i] * cols, - fDatasets[fMinor].GetData() + (fSamples[i] + 1) * cols, - OversampledMinorTensor.GetData() + index * cols); - index++; - } + GatherSamples(ShuffledTensor, fDatasets[fMinor], fNumResampledMinor, fDatasets[fMajor]); +} - fTensorOperators->ConcatenateTensors(SampledTensor, {OversampledMinorTensor, fDatasets[fMajor]}); - if (fShuffle) { - fTensorOperators->ShuffleTensor(ShuffledTensor, SampledTensor); - } else { - ShuffledTensor = std::move(SampledTensor); +////////////////////////////////////////////////////////////////////////// +/// \brief Fill the sampled entries: the resampled entries followed by all the other ones, in shuffled order +/// \param[in] ShuffledTensor Tensor with all the sampled entries +/// \param[in] ResampledDataset Dataset the entries in fSamples are taken from +/// \param[in] numResampled Number of resampled entries +/// \param[in] OtherDataset Dataset whose entries are all taken +void RSampler::GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix &ResampledDataset, + std::size_t numResampled, const RFlat2DMatrix &OtherDataset) +{ + std::size_t cols = fDatasets[0].GetCols(); + const auto permutation = fTensorOperators->MakePermutation(fNumEntries); + + // read the entries straight from the datasets, instead of concatenating and shuffling copies of them + ShuffledTensor.Resize(fNumEntries, cols); + float *out = ShuffledTensor.GetData(); + for (std::size_t i = 0; i < fNumEntries; i++) { + const std::size_t entry = permutation.empty() ? i : permutation[i]; + const float *row = entry < numResampled ? ResampledDataset.GetData() + fSamples[entry] * cols + : OtherDataset.GetData() + (entry - numResampled) * cols; + out = std::copy_n(row, cols, out); } } diff --git a/tree/ml/src/RSampler.hxx b/tree/ml/src/RSampler.hxx index f9a795e12a21e..7f35298df31ea 100644 --- a/tree/ml/src/RSampler.hxx +++ b/tree/ml/src/RSampler.hxx @@ -68,6 +68,9 @@ public: void RandomOversampler(RFlat2DMatrix &ShuffledTensor); + void GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix &ResampledDataset, std::size_t numResampled, + const RFlat2DMatrix &OtherDataset); + void SampleWithReplacement(std::size_t n_samples, std::size_t max); void SampleWithoutReplacement(std::size_t n_samples, std::size_t max); From 6ef3b1b2613aacdd5a9c4eb4174062ee27d5c494 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 11:24:34 +0200 Subject: [PATCH 3/8] [ML] Opt.: Load the filtered eager dataset in a single event loop The filtered eager path ran one event loop per passing entry, after a `Take` of `rdfentry_` to find them. A single event loop over the filtered dataframe loads the same entries, and the buffer is reserved from the cluster metadata like in the unfiltered case. Assisted-by: Claude Opus 5.5 --- tree/ml/src/RDatasetLoader.cxx | 49 +++++++++++----------------------- 1 file changed, 16 insertions(+), 33 deletions(-) diff --git a/tree/ml/src/RDatasetLoader.cxx b/tree/ml/src/RDatasetLoader.cxx index 4d280be9bb55f..6202cd7f02622 100644 --- a/tree/ml/src/RDatasetLoader.cxx +++ b/tree/ml/src/RDatasetLoader.cxx @@ -28,49 +28,32 @@ RFlat2DMatrix RDatasetLoader::LoadDataframe(ROOT::RDF::RNode &rdf) { const bool NotFiltered = rdf.GetFilterNames().empty(); - // size the buffer from the cluster metadata, Count() is only the fallback for sources without it - ROOT::RDF::RResultPtr> Entries; + // size the buffer from the cluster metadata, Count() is only the fallback for sources without it. With filters + // this is an upper bound: the memory reserved for the entries that don't pass them is never touched std::size_t NumEntries = 0; - if (NotFiltered) { - try { - for (const auto &cluster : ROOT::Internal::RDF::GetDatasetGlobalClusterBoundaries(rdf)) { - NumEntries += cluster.second - cluster.first; - } - } catch (const std::runtime_error &) { - // GetDatasetGlobalClusterBoundaries() throws when the RDataFrame has no cluster - // metadata to query (a source other than TTree/RNTuple or no data source at all). - // Those paths fall back to Count() below. + try { + for (const auto &cluster : ROOT::Internal::RDF::GetDatasetGlobalClusterBoundaries(rdf)) { + NumEntries += cluster.second - cluster.first; } - if (NumEntries == 0) { - NumEntries = static_cast(*rdf.Count()); - } - } else { - Entries = rdf.Take("rdfentry_"); - NumEntries = Entries->size(); - // add the last element in entries to not go out of range when filling chunks - Entries->push_back((*Entries)[NumEntries - 1] + 1); + } catch (const std::runtime_error &) { + // GetDatasetGlobalClusterBoundaries() throws when the RDataFrame has no cluster + // metadata to query (a source other than TTree/RNTuple or no data source at all). + // Those paths fall back to Count() below. + } + if (NumEntries == 0) { + NumEntries = static_cast(*rdf.Count()); } // reserve the dataset and let the values be appended to it RFlat2DMatrix Dataset; Dataset.fRVec.reserve(NumEntries * fNumDatasetCols); - if (NotFiltered) { - ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, Dataset.fRVec); - } - - else { - for (std::size_t j = 0; j < NumEntries; j++) { - ROOT::Internal::RDF::ChangeBeginAndEndEntries(rdf, (*Entries)[j], (*Entries)[j + 1]); - ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, Dataset.fRVec); - } - } - Dataset.Resize(NumEntries, fNumDatasetCols); - - // reset dataframe, only the filtered path changed the entry range + // a single event loop, the filters only let the passing entries through + ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, Dataset.fRVec); if (!NotFiltered) { - ROOT::Internal::RDF::ChangeBeginAndEndEntries(rdf, (*Entries)[0], (*Entries)[NumEntries]); + NumEntries = Dataset.GetSize() / fNumDatasetCols; } + Dataset.Resize(NumEntries, fNumDatasetCols); return Dataset; } From ada21d65943f3d41db16be185e6be98785164412 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 11:24:40 +0200 Subject: [PATCH 4/8] [ML] Opt.: Load the filtered clusters once in the first lazy epoch The first epoch read each filtered training cluster twice: once to collect `rdfentry_` and find the training/validation boundary, then again to load the training rows. The entries are now taken in the same event loop that loads the cluster, and the validation rows are dropped from the buffer. Assisted-by: Claude Opus 5.5 --- tree/ml/src/RClusterLoader.cxx | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tree/ml/src/RClusterLoader.cxx b/tree/ml/src/RClusterLoader.cxx index 0eb36664c46f6..b5fcd9accbcd1 100644 --- a/tree/ml/src/RClusterLoader.cxx +++ b/tree/ml/src/RClusterLoader.cxx @@ -205,13 +205,14 @@ std::size_t RClusterLoader::LoadTrainingClusterInto(RFlat2DMatrix &dest, std::si } ROOT::RDF::RNode &rdf = fRdfs[rdfIdx]; + const std::size_t numLoadedValues = dest.fRVec.size(); - std::vector rdfEntries; - rdfEntries.reserve(endRow - startRow); - + // load the whole cluster and capture rdfentry_ in the same event loop ROOT::Internal::RDF::ChangeBeginAndEndEntries(rdf, startRow, endRow); - rdf.Foreach([&](ULong64_t entry) { rdfEntries.push_back(entry); }, {"rdfentry_"}); + auto entries = rdf.Take("rdfentry_"); + ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, dest.fRVec); ROOT::Internal::RDF::ChangeBeginAndEndEntries(rdf, 0, fRdfSizes[rdfIdx]); + auto &rdfEntries = *entries; const std::size_t totalFiltered = rdfEntries.size(); if (totalFiltered == 0) { @@ -262,8 +263,11 @@ std::size_t RClusterLoader::LoadTrainingClusterInto(RFlat2DMatrix &dest, std::si fAccumulatedFilteredForTrain += trainCount; fAccumulatedFilteredForVal += valCount; - if (trainCount > 0) - LoadClusterInto(dest, rdfIdx, trainStart, trainEnd); + // keep only the training rows of the cluster, the validation ones are loaded by the validation epoch + float *rows = dest.fRVec.data() + numLoadedValues; + if (!trainIsPrefix) + std::copy(rows + valCount * fNumChunkCols, rows + totalFiltered * fNumChunkCols, rows); + dest.fRVec.resize(numLoadedValues + trainCount * fNumChunkCols); return trainCount; } From f15bb7bb21b5f9231a9e898ccf5b0dfc4f4da252 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 11:25:28 +0200 Subject: [PATCH 5/8] [ML] Opt.: Seed the sampler's random generator once per sampling `SampleWithReplacement` created and seeded a new generator for every sample, which dominated the oversampling epochs and, with a fixed seed, drew the same sample every time. Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 31 +++++++++++++++++++ tree/ml/src/RSampler.cxx | 20 ++++++------ 2 files changed, 42 insertions(+), 9 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index c1682065b0a89..443404c1c0a94 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -7838,6 +7838,37 @@ def test16_vector_padding(self): self.teardown_file(self.file_name5) raise + def test17_set_seed_resamples_different_entries(self): + # with a fixed seed, the resampled minority entries used to all be the same one + self.create_file_major() + self.create_file_minor() + + try: + df_major = ROOT.RDataFrame(self.tree_name, self.file_name1) + df_minor = ROOT.RDataFrame(self.tree_name, self.file_name2) + + dl = ROOT.Experimental.ML.RDataLoader( + [df_major, df_minor], + batch_size=2, + target="b2", + shuffle=True, + drop_remainder=False, + set_seed=42, + load_eager=True, + sampling_type="oversampling", + sampling_ratio=1.0, + ) + + x = np.concatenate([x for x, _ in dl.as_numpy()]).flatten() + minor = x[x % 2 == 1] + + self.assertEqual(len(minor), 10) + self.assertGreater(len(np.unique(minor)), 1) + + finally: + self.teardown_file(self.file_name1) + self.teardown_file(self.file_name2) + class DataLoaderSave(unittest.TestCase): file_name = "dataloader_save_input.root" diff --git a/tree/ml/src/RSampler.cxx b/tree/ml/src/RSampler.cxx index 6285d48b13996..de004baf763a6 100644 --- a/tree/ml/src/RSampler.cxx +++ b/tree/ml/src/RSampler.cxx @@ -150,20 +150,22 @@ void RSampler::GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix void RSampler::SampleWithReplacement(std::size_t n_samples, std::size_t max) { std::uniform_int_distribution<> dist(0, max - 1); + // a single generator for all the samples: seeding one per sample was slow, and with a fixed seed it drew the same + // sample every time + std::random_device rd; + std::mt19937 g; + + if (fSetSeed == 0) { + g.seed(rd()); + } else { + g.seed(fSetSeed); + } + fSamples.clear(); fSamples.reserve(n_samples); for (std::size_t i = 0; i < n_samples; ++i) { std::size_t sample; if (fShuffle) { - std::random_device rd; - std::mt19937 g; - - if (fSetSeed == 0) { - g.seed(rd()); - } else { - g.seed(fSetSeed); - } - sample = dist(g); } From ab22a7db739f2ed3d216a4c5db7b9c883e0e1ff0 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 11:26:40 +0200 Subject: [PATCH 6/8] [ML] Opt.: Create the eager batches when they are requested The eager epochs copied the whole dataset into batches before returning the first one, so the dataset and all its batches were in memory together. Each batch is now gathered from the dataset when it is requested. Assisted-by: Claude Opus 5.5 --- tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx | 12 ++++++ tree/ml/src/RBatchLoader.cxx | 19 ++++++---- tree/ml/src/RBatchLoader.hxx | 6 ++- tree/ml/src/RDataLoaderEngine.cxx | 46 +++++++++++++++++++---- 4 files changed, 65 insertions(+), 18 deletions(-) diff --git a/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx b/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx index 7459a464c34ed..d0bf164ec8e5d 100644 --- a/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx +++ b/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx @@ -121,6 +121,18 @@ private: std::size_t fTrainingEpochCount{0}; std::size_t fValidationEpochCount{0}; + /// \brief Rows of an eager epoch, turned into batches as they are requested + struct REagerEpoch { + const RFlat2DMatrix *fSource = nullptr; + std::vector fPermutation; + std::size_t fNextRow = 0; + }; + + REagerEpoch fTrainingEagerEpoch; + REagerEpoch fValidationEagerEpoch; + + void CreateNextEagerBatch(RBatchLoader &batchLoader, REagerEpoch &epoch); + std::vector MakeColumnLayout(); /// \brief Opens a training or validation epoch and closes it again when done diff --git a/tree/ml/src/RBatchLoader.cxx b/tree/ml/src/RBatchLoader.cxx index a6caef30c34c7..89d84ab7304c3 100644 --- a/tree/ml/src/RBatchLoader.cxx +++ b/tree/ml/src/RBatchLoader.cxx @@ -84,14 +84,15 @@ void RBatchLoader::MarkProducerDone() /// After the batch has been processed, it should be destroyed. /// \param[in] chunkTensor Tensor with the data from the chunk /// \param[in] permutation Order in which to read the rows of the chunk, empty to read them in order -/// \param[in] idxs Index of batch in the chunk +/// \param[in] firstRow First row of the batch in the chunk /// \return Batch std::unique_ptr RBatchLoader::CreateBatch(const RFlat2DMatrix &chunkTensor, - const std::vector &permutation, std::size_t idxs) + const std::vector &permutation, + std::size_t firstRow) { auto batch = std::make_unique(); batch->Reserve(fBatchSize, fNumColumns); - RFlat2DMatrixOperators::AppendRows(*batch, chunkTensor, permutation, idxs * fBatchSize, (idxs + 1) * fBatchSize); + RFlat2DMatrixOperators::AppendRows(*batch, chunkTensor, permutation, firstRow, firstRow + fBatchSize); return batch; } @@ -126,11 +127,13 @@ RFlat2DMatrix RBatchLoader::GetBatch() /// \brief Creating the batches from a chunk and add them to the queue. /// \param[in] chunkTensor Tensor with the data from the chunk /// \param[in] permutation Order in which to read the rows of the chunk, empty to read them in order +/// \param[in] beginRow First row of the chunk to create batches from +/// \param[in] endRow End of the rows of the chunk to create batches from /// \param[in] isLastBatch Check if the batch in the chunk is the last one void RBatchLoader::CreateBatches(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, - bool isLastBatch) + std::size_t beginRow, std::size_t endRow, bool isLastBatch) { - std::size_t chunkSize = chunkTensor.GetRows(); + std::size_t chunkSize = endRow - beginRow; std::size_t numBatches = chunkSize / fBatchSize; // create a vector of batches @@ -138,12 +141,12 @@ void RBatchLoader::CreateBatches(const RFlat2DMatrix &chunkTensor, const std::ve // fill the full batches from the chunk into a vector for (std::size_t i = 0; i < numBatches; i++) { - batches.emplace_back(CreateBatch(chunkTensor, permutation, i)); + batches.emplace_back(CreateBatch(chunkTensor, permutation, beginRow + i * fBatchSize)); } // append the remaining entries of the chunk to fPrimaryLeftoverBatch, and add it to the batches once it is full - for (std::size_t row = numBatches * fBatchSize; row < chunkSize;) { - const std::size_t numRows = std::min(chunkSize - row, fBatchSize - fPrimaryLeftoverBatch->GetRows()); + for (std::size_t row = beginRow + numBatches * fBatchSize; row < endRow;) { + const std::size_t numRows = std::min(endRow - row, fBatchSize - fPrimaryLeftoverBatch->GetRows()); fPrimaryLeftoverBatch->Reserve(fBatchSize, fNumColumns); RFlat2DMatrixOperators::AppendRows(*fPrimaryLeftoverBatch, chunkTensor, permutation, row, row + numRows); row += numRows; diff --git a/tree/ml/src/RBatchLoader.hxx b/tree/ml/src/RBatchLoader.hxx index 70ad48cdd7532..8159a82f555e7 100644 --- a/tree/ml/src/RBatchLoader.hxx +++ b/tree/ml/src/RBatchLoader.hxx @@ -75,12 +75,14 @@ public: void MarkProducerDone(); std::unique_ptr - CreateBatch(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, std::size_t idxs); + CreateBatch(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, std::size_t firstRow); RFlat2DMatrix GetBatch(); - void CreateBatches(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, bool isLastBatch); + void CreateBatches(const RFlat2DMatrix &chunkTensor, const std::vector &permutation, + std::size_t beginRow, std::size_t endRow, bool isLastBatch); void RecalculateBatchCounts(std::size_t numEntries); bool isProducerDone() { return fProducerDone; } + std::size_t GetBatchSize() { return fBatchSize; } std::size_t GetNumBatches() { return fNumBatches; } std::size_t GetNumEntries() { return fNumEntries; } std::size_t GetNumRemainderRows() { return fLeftoverBatchSize; } diff --git a/tree/ml/src/RDataLoaderEngine.cxx b/tree/ml/src/RDataLoaderEngine.cxx index 37b61fc13287d..d051b400b4c5d 100644 --- a/tree/ml/src/RDataLoaderEngine.cxx +++ b/tree/ml/src/RDataLoaderEngine.cxx @@ -356,8 +356,9 @@ void RDataLoaderEngine::LoadData() stagingBuffer.Resize(rowOffset, fClusterLoader->GetNumChunkCols()); // the batches read the buffer in shuffled order, so it needs no shuffled copy - fTrainingBatchLoader->CreateBatches( - stagingBuffer, fTensorOperators->MakePermutation(stagingBuffer.GetRows()), isLastBuffer); + fTrainingBatchLoader->CreateBatches(stagingBuffer, + fTensorOperators->MakePermutation(stagingBuffer.GetRows()), 0, + stagingBuffer.GetRows(), isLastBuffer); // Re-acquire the lock before the next iteration to check conditions and update indices lock.lock(); @@ -415,8 +416,9 @@ void RDataLoaderEngine::LoadData() } stagingBuffer.Resize(rowOffset, fClusterLoader->GetNumChunkCols()); - fValidationBatchLoader->CreateBatches( - stagingBuffer, fTensorOperators->MakePermutation(stagingBuffer.GetRows()), isLastBuffer); + fValidationBatchLoader->CreateBatches(stagingBuffer, + fTensorOperators->MakePermutation(stagingBuffer.GetRows()), 0, + stagingBuffer.GetRows(), isLastBuffer); lock.lock(); } @@ -443,8 +445,8 @@ void RDataLoaderEngine::CreateTrainBatches() fTrainingSampler->Sampler(fSampledTrainingDataset); } - fTrainingBatchLoader->CreateBatches(*source, permutation, true); - fTrainingBatchLoader->MarkProducerDone(); + // the batches are created one at a time when they are requested, see GetTrainBatch() + fTrainingEagerEpoch = {source, std::move(permutation), 0}; } } @@ -466,14 +468,38 @@ void RDataLoaderEngine::CreateValidationBatches() fValidationSampler->Sampler(fSampledValidationDataset); } - fValidationBatchLoader->CreateBatches(*source, permutation, true); - fValidationBatchLoader->MarkProducerDone(); + // the batches are created one at a time when they are requested, see GetValidationBatch() + fValidationEagerEpoch = {source, std::move(permutation), 0}; + } +} + +/// \brief Create the next batch of an eager epoch, unless one is already queued or the epoch is done +void RDataLoaderEngine::CreateNextEagerBatch(RBatchLoader &batchLoader, REagerEpoch &epoch) +{ + if (!epoch.fSource) { + return; + } + + const std::size_t numRows = epoch.fSource->GetRows(); + // a slice of one batch size gives a full batch, the last slice also gives the leftover batch + while (batchLoader.GetNumBatchQueue() == 0 && !batchLoader.isProducerDone()) { + const std::size_t endRow = + std::min(epoch.fNextRow + std::max(batchLoader.GetBatchSize(), 1), numRows); + batchLoader.CreateBatches(*epoch.fSource, epoch.fPermutation, epoch.fNextRow, endRow, endRow == numRows); + epoch.fNextRow = endRow; + if (endRow == numRows) { + batchLoader.MarkProducerDone(); + } } } /// \brief Loads a training batch from the queue RFlat2DMatrix RDataLoaderEngine::GetTrainBatch() { + if (fLoadEager) { + CreateNextEagerBatch(*fTrainingBatchLoader, fTrainingEagerEpoch); + } + // Get next batch if available return fTrainingBatchLoader->GetBatch(); } @@ -481,6 +507,10 @@ RFlat2DMatrix RDataLoaderEngine::GetTrainBatch() /// \brief Loads a validation batch from the queue RFlat2DMatrix RDataLoaderEngine::GetValidationBatch() { + if (fLoadEager) { + CreateNextEagerBatch(*fValidationBatchLoader, fValidationEagerEpoch); + } + // Get next batch if available return fValidationBatchLoader->GetBatch(); } From 2aeb6cf5da6cfef7af8711889059c94ffcb3edeb Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 11:51:31 +0200 Subject: [PATCH 7/8] [ML] Opt.: Split the eager dataset by row indices instead of copying it The dataframes are loaded into a single dataset, and the training and validation sets, as well as the classes of the samplers, are lists of its rows. The split no longer copies the dataset, nor do the samplers, and the batches gather the rows of their set. Assisted-by: Claude Opus 5.5 --- .../pythonizations/test/ml_dataloader.py | 17 ++ tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx | 13 +- tree/ml/src/RDataLoaderEngine.cxx | 60 +++---- tree/ml/src/RDatasetLoader.cxx | 157 ++++++++---------- tree/ml/src/RDatasetLoader.hxx | 22 +-- tree/ml/src/RFlat2DMatrixOperators.cxx | 23 ++- tree/ml/src/RFlat2DMatrixOperators.hxx | 3 + tree/ml/src/RSampler.cxx | 39 ++--- tree/ml/src/RSampler.hxx | 12 +- 9 files changed, 189 insertions(+), 157 deletions(-) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 443404c1c0a94..7d0cf6bd07f5f 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -3606,6 +3606,23 @@ def test16_vector_padding(self): self.teardown_file(self.file_name5) raise + def test17_dataframes_without_cluster_metadata(self): + # the dataframes are not read from a file, so their entries are counted with Count() + dl = ROOT.Experimental.ML.RDataLoader( + [self.define_rdf1(), self.define_rdf2()], + batch_size=3, + target="b2", + shuffle=False, + drop_remainder=False, + load_eager=True, + ) + gen_train, gen_validation = dl.train_test_split(0.4) + + x_train = [v for x, _ in gen_train.as_numpy() for v in x.flatten().tolist()] + x_val = [v for x, _ in gen_validation.as_numpy() for v in x.flatten().tolist()] + self.assertEqual(x_train, [0.0, 1.0, 2.0, 5.0, 6.0, 7.0]) + self.assertEqual(x_val, [3.0, 4.0, 8.0, 9.0]) + class DataLoaderLazyLoadingMultipleDataframes(unittest.TestCase): file_name1 = "first_half.root" diff --git a/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx b/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx index d0bf164ec8e5d..5c2bf4edca66d 100644 --- a/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx +++ b/tree/ml/inc/ROOT/ML/RDataLoaderEngine.hxx @@ -109,11 +109,11 @@ private: std::size_t fNumValidationEntries; // flattened buffers for chunks and temporary tensors (rows * cols) - std::vector fTrainingDatasets; - std::vector fValidationDatasets; + RFlat2DMatrix fDataset; - RFlat2DMatrix fTrainingDataset; - RFlat2DMatrix fValidationDataset; + // rows of fDataset in the training and validation sets, one set per dataframe when sampling + std::vector> fTrainingRows; + std::vector> fValidationRows; RFlat2DMatrix fSampledTrainingDataset; RFlat2DMatrix fSampledValidationDataset; @@ -123,14 +123,19 @@ private: /// \brief Rows of an eager epoch, turned into batches as they are requested struct REagerEpoch { + // non-owning pointer: the data stays in the engine's own members (fDataset for the normal path, + // fSampledTraining/ValidationDataset for the sampler path) const RFlat2DMatrix *fSource = nullptr; std::vector fPermutation; std::size_t fNextRow = 0; + std::size_t fEndRow = 0; }; REagerEpoch fTrainingEagerEpoch; REagerEpoch fValidationEagerEpoch; + REagerEpoch MakeEagerEpoch(const std::vector &rows); + void CreateNextEagerBatch(RBatchLoader &batchLoader, REagerEpoch &epoch); std::vector MakeColumnLayout(); diff --git a/tree/ml/src/RDataLoaderEngine.cxx b/tree/ml/src/RDataLoaderEngine.cxx index d051b400b4c5d..4436945ab82bb 100644 --- a/tree/ml/src/RDataLoaderEngine.cxx +++ b/tree/ml/src/RDataLoaderEngine.cxx @@ -78,23 +78,20 @@ RDataLoaderEngine::RDataLoaderEngine(const std::vector &rdfs, fDatasetLoader = std::make_unique(fRdfs, fTestSize, fCols, fVecSizes, vecPadding, fShuffle, fSetSeed); fDatasetLoader->SplitDatasets(fSampleType == ""); + fDataset = fDatasetLoader->ReleaseDataset(); + fTrainingRows = fDatasetLoader->ReleaseTrainingRows(); + fValidationRows = fDatasetLoader->ReleaseValidationRows(); if (fSampleType == "") { - fTrainingDataset = fDatasetLoader->ReleaseTrainingDataset(); - fValidationDataset = fDatasetLoader->ReleaseValidationDataset(); - fNumTrainingEntries = fDatasetLoader->GetNumTrainingEntries(); fNumValidationEntries = fDatasetLoader->GetNumValidationEntries(); } else { - fTrainingDatasets = fDatasetLoader->ReleaseTrainingDatasets(); - fValidationDatasets = fDatasetLoader->ReleaseValidationDatasets(); - - fTrainingSampler = - std::make_unique(fTrainingDatasets, fSampleType, fSampleRatio, fReplacement, fShuffle, fSetSeed); - fValidationSampler = std::make_unique(fValidationDatasets, fSampleType, fSampleRatio, fReplacement, - fShuffle, fSetSeed); + fTrainingSampler = std::make_unique(fDataset, fTrainingRows, fSampleType, fSampleRatio, fReplacement, + fShuffle, fSetSeed); + fValidationSampler = std::make_unique(fDataset, fValidationRows, fSampleType, fSampleRatio, + fReplacement, fShuffle, fSetSeed); fNumTrainingEntries = fTrainingSampler->GetNumEntries(); fNumValidationEntries = fValidationSampler->GetNumEntries(); @@ -433,20 +430,16 @@ void RDataLoaderEngine::CreateTrainBatches() fTrainingBatchLoader->Activate(); if (fLoadEager) { - RFlat2DMatrix *source = &fSampledTrainingDataset; - std::vector permutation; + // the batches are created one at a time when they are requested, see GetTrainBatch() if (fSampleType == "") { - // the batches read the dataset in shuffled order, so it needs no shuffled copy - source = &fTrainingDataset; - permutation = fTensorOperators->MakePermutation(fTrainingDataset.GetRows()); + // the batches read the training rows of the dataset in shuffled order, so they need no copy + fTrainingEagerEpoch = MakeEagerEpoch(fTrainingRows[0]); } else { fTrainingSampler->Sampler(fSampledTrainingDataset); + fTrainingEagerEpoch = {&fSampledTrainingDataset, {}, 0, fSampledTrainingDataset.GetRows()}; } - - // the batches are created one at a time when they are requested, see GetTrainBatch() - fTrainingEagerEpoch = {source, std::move(permutation), 0}; } } @@ -457,20 +450,32 @@ void RDataLoaderEngine::CreateValidationBatches() fValidationBatchLoader->Activate(); if (fLoadEager) { - RFlat2DMatrix *source = &fSampledValidationDataset; - std::vector permutation; + // the batches are created one at a time when they are requested, see GetValidationBatch() if (fSampleType == "") { - source = &fValidationDataset; - permutation = fTensorOperators->MakePermutation(fValidationDataset.GetRows()); + fValidationEagerEpoch = MakeEagerEpoch(fValidationRows[0]); } else { fValidationSampler->Sampler(fSampledValidationDataset); + fValidationEagerEpoch = {&fSampledValidationDataset, {}, 0, fSampledValidationDataset.GetRows()}; } + } +} - // the batches are created one at a time when they are requested, see GetValidationBatch() - fValidationEagerEpoch = {source, std::move(permutation), 0}; +/// \brief The eager epoch reading the given rows of the dataset, in shuffled order when shuffling +RDataLoaderEngine::REagerEpoch RDataLoaderEngine::MakeEagerEpoch(const std::vector &rows) +{ + // contiguous rows read in order are a range of the dataset, they need no permutation + if (!fShuffle && (rows.empty() || rows.back() - rows.front() + 1 == rows.size())) { + const std::size_t beginRow = rows.empty() ? 0 : rows.front(); + return {&fDataset, {}, beginRow, beginRow + rows.size()}; } + + // shuffling a copy of the rows gives the same order as reading them through a shuffled permutation + std::vector permutation = rows; + fTensorOperators->ShuffleIndices(permutation); + const std::size_t numRows = permutation.size(); + return {&fDataset, std::move(permutation), 0, numRows}; } /// \brief Create the next batch of an eager epoch, unless one is already queued or the epoch is done @@ -480,14 +485,13 @@ void RDataLoaderEngine::CreateNextEagerBatch(RBatchLoader &batchLoader, REagerEp return; } - const std::size_t numRows = epoch.fSource->GetRows(); // a slice of one batch size gives a full batch, the last slice also gives the leftover batch while (batchLoader.GetNumBatchQueue() == 0 && !batchLoader.isProducerDone()) { const std::size_t endRow = - std::min(epoch.fNextRow + std::max(batchLoader.GetBatchSize(), 1), numRows); - batchLoader.CreateBatches(*epoch.fSource, epoch.fPermutation, epoch.fNextRow, endRow, endRow == numRows); + std::min(epoch.fNextRow + std::max(batchLoader.GetBatchSize(), 1), epoch.fEndRow); + batchLoader.CreateBatches(*epoch.fSource, epoch.fPermutation, epoch.fNextRow, endRow, endRow == epoch.fEndRow); epoch.fNextRow = endRow; - if (endRow == numRows) { + if (endRow == epoch.fEndRow) { batchLoader.MarkProducerDone(); } } diff --git a/tree/ml/src/RDatasetLoader.cxx b/tree/ml/src/RDatasetLoader.cxx index 6202cd7f02622..423993bdd431a 100644 --- a/tree/ml/src/RDatasetLoader.cxx +++ b/tree/ml/src/RDatasetLoader.cxx @@ -1,5 +1,7 @@ #include "RDatasetLoader.hxx" +#include "ROOT/RDFHelpers.hxx" + namespace ROOT::Experimental::Internal::ML { RDatasetLoader::RDatasetLoader(const std::vector &rdfs, const float validationSplit, @@ -21,119 +23,106 @@ RDatasetLoader::RDatasetLoader(const std::vector &rdfs, const } ////////////////////////////////////////////////////////////////////////// -/// \brief Load all the entries of an individual dataframe -/// \param[in] rdf Dataframe to load -/// \return Tensor with the dataset -RFlat2DMatrix RDatasetLoader::LoadDataframe(ROOT::RDF::RNode &rdf) +/// \brief Number of entries to reserve for each dataframe +/// \param[in] rdfs Dataframes to load +/// \return Number of entries of each dataframe, an upper bound for a filtered dataframe +std::vector RDatasetLoader::CountEntries(std::vector &rdfs) { - const bool NotFiltered = rdf.GetFilterNames().empty(); - // size the buffer from the cluster metadata, Count() is only the fallback for sources without it. With filters // this is an upper bound: the memory reserved for the entries that don't pass them is never touched - std::size_t NumEntries = 0; - try { - for (const auto &cluster : ROOT::Internal::RDF::GetDatasetGlobalClusterBoundaries(rdf)) { - NumEntries += cluster.second - cluster.first; + std::vector NumEntries(rdfs.size(), 0); + std::vector> Counts; + for (std::size_t i = 0; i < rdfs.size(); i++) { + try { + for (const auto &cluster : ROOT::Internal::RDF::GetDatasetGlobalClusterBoundaries(rdfs[i])) { + NumEntries[i] += cluster.second - cluster.first; + } + } catch (const std::runtime_error &) { + // GetDatasetGlobalClusterBoundaries() throws when the RDataFrame has no cluster + // metadata to query (a source other than TTree/RNTuple or no data source at all). + // Those paths fall back to Count() below. } - } catch (const std::runtime_error &) { - // GetDatasetGlobalClusterBoundaries() throws when the RDataFrame has no cluster - // metadata to query (a source other than TTree/RNTuple or no data source at all). - // Those paths fall back to Count() below. + if (NumEntries[i] == 0) { + Counts.push_back(rdfs[i].Count()); + } + } + + // the dataframes without cluster metadata are counted together + if (!Counts.empty()) { + ROOT::RDF::RunGraphs({Counts.begin(), Counts.end()}); } - if (NumEntries == 0) { - NumEntries = static_cast(*rdf.Count()); + for (std::size_t i = 0, c = 0; i < rdfs.size(); i++) { + if (NumEntries[i] == 0) { + NumEntries[i] = static_cast(*Counts[c++]); + } } - // reserve the dataset and let the values be appended to it - RFlat2DMatrix Dataset; - Dataset.fRVec.reserve(NumEntries * fNumDatasetCols); + return NumEntries; +} + +////////////////////////////////////////////////////////////////////////// +/// \brief Load all the entries of an individual dataframe at the end of the dataset +/// \param[in] rdf Dataframe to load +/// \param[in] NumEntries Number of entries of the dataframe, an upper bound for a filtered dataframe +void RDatasetLoader::LoadDataframe(ROOT::RDF::RNode &rdf, std::size_t NumEntries) +{ + const bool NotFiltered = rdf.GetFilterNames().empty(); + const std::size_t firstRow = fDataset.GetRows(); // a single event loop, the filters only let the passing entries through - ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, Dataset.fRVec); + ROOT::Internal::RDF::LoadCustomValues(rdf, fCols, fVecSizes, fVecPadding, fDataset.fRVec); if (!NotFiltered) { - NumEntries = Dataset.GetSize() / fNumDatasetCols; + NumEntries = fDataset.GetSize() / fNumDatasetCols - firstRow; } - Dataset.Resize(NumEntries, fNumDatasetCols); - - return Dataset; + fDataset.Resize(firstRow + NumEntries, fNumDatasetCols); } ////////////////////////////////////////////////////////////////////////// -/// \brief Split a loaded dataset into training and validation entries, appended to the given datasets -/// \param[in] Dataset Tensor with the loaded dataset -/// \param[in] TrainingDataset Tensor for the training dataset -/// \param[in] ValidationDataset Tensor for the validation dataset -void RDatasetLoader::SplitDataset(RFlat2DMatrix &Dataset, RFlat2DMatrix &TrainingDataset, - RFlat2DMatrix &ValidationDataset) +/// \brief Split the rows of the dataset from \p firstRow on into training and validation rows +/// \param[in] firstRow First row of the dataframe in the dataset +/// \param[in] TrainingRows Rows of the training set, appended to +/// \param[in] ValidationRows Rows of the validation set, appended to +void RDatasetLoader::SplitRows(std::size_t firstRow, std::vector &TrainingRows, + std::vector &ValidationRows) { // number of training and validation entries after the split - std::size_t NumEntries = Dataset.GetRows(); + std::size_t NumEntries = fDataset.GetRows() - firstRow; std::size_t NumValidationEntries = static_cast(fValidationSplit * NumEntries); std::size_t NumTrainingEntries = NumEntries - NumValidationEntries; - // read the entries in shuffled order, instead of copying the dataset into a shuffled one first + // the sets refer to the rows of the dataset in shuffled order, instead of copying them const auto permutation = fTensorOperators->MakePermutation(NumEntries); - - ValidationDataset.fRVec.reserve(ValidationDataset.GetSize() + NumValidationEntries * fNumDatasetCols); - RFlat2DMatrixOperators::AppendRows(ValidationDataset, Dataset, permutation, NumTrainingEntries, NumEntries); - - TrainingDataset.fRVec.reserve(TrainingDataset.GetSize() + NumTrainingEntries * fNumDatasetCols); - RFlat2DMatrixOperators::AppendRows(TrainingDataset, Dataset, permutation, 0, NumTrainingEntries); + TrainingRows.reserve(TrainingRows.size() + NumTrainingEntries); + ValidationRows.reserve(ValidationRows.size() + NumValidationEntries); + for (std::size_t i = 0; i < NumEntries; i++) { + const std::size_t row = firstRow + (permutation.empty() ? i : permutation[i]); + (i < NumTrainingEntries ? TrainingRows : ValidationRows).push_back(row); + } } ////////////////////////////////////////////////////////////////////////// -/// \brief Split the dataframes in a training and validation dataset -/// \param[in] concatenate Whether to concatenate the datasets of all the dataframes into one +/// \brief Load the dataframes into a single dataset and split it in training and validation rows +/// \param[in] concatenate Whether to concatenate the rows of all the dataframes into a single set void RDatasetLoader::SplitDatasets(bool concatenate) { - std::vector Datasets; - for (auto &rdf : f_rdfs) { - Datasets.push_back(LoadDataframe(rdf)); - } - - if (concatenate && Datasets.size() > 1) { - // split every dataset directly into the concatenated ones, sized for all of them - std::size_t NumValidationEntries = 0; - std::size_t NumEntries = 0; - for (const auto &Dataset : Datasets) { - NumEntries += Dataset.GetRows(); - NumValidationEntries += static_cast(fValidationSplit * Dataset.GetRows()); - } - fTrainingDataset.fRVec.reserve((NumEntries - NumValidationEntries) * fNumDatasetCols); - fValidationDataset.fRVec.reserve(NumValidationEntries * fNumDatasetCols); - - for (auto &Dataset : Datasets) { - SplitDataset(Dataset, fTrainingDataset, fValidationDataset); - Dataset = RFlat2DMatrix(); - } - } else { - for (auto &Dataset : Datasets) { - RFlat2DMatrix TrainingDataset; - RFlat2DMatrix ValidationDataset; - SplitDataset(Dataset, TrainingDataset, ValidationDataset); - Dataset = RFlat2DMatrix(); - fTrainingDatasets.push_back(std::move(TrainingDataset)); - fValidationDatasets.push_back(std::move(ValidationDataset)); - } - - if (concatenate) { - fTrainingDataset = std::move(fTrainingDatasets[0]); - fValidationDataset = std::move(fValidationDatasets[0]); - fTrainingDatasets.clear(); - fValidationDatasets.clear(); - } + // reserve the dataset for all the dataframes and let their values be appended to it + const std::vector NumEntries = CountEntries(f_rdfs); + fDataset.Reserve(std::accumulate(NumEntries.begin(), NumEntries.end(), std::size_t{0}), fNumDatasetCols); + + const std::size_t numSets = concatenate ? 1 : f_rdfs.size(); + fTrainingRows.resize(numSets); + fValidationRows.resize(numSets); + for (std::size_t i = 0; i < f_rdfs.size(); i++) { + const std::size_t firstRow = fDataset.GetRows(); + LoadDataframe(f_rdfs[i], NumEntries[i]); + SplitRows(firstRow, fTrainingRows[concatenate ? 0 : i], fValidationRows[concatenate ? 0 : i]); } fNumTrainingEntries = 0; fNumValidationEntries = 0; - if (concatenate) { - fNumTrainingEntries = fTrainingDataset.GetRows(); - fNumValidationEntries = fValidationDataset.GetRows(); - } else { - for (std::size_t i = 0; i < fTrainingDatasets.size(); i++) { - fNumTrainingEntries += fTrainingDatasets[i].GetRows(); - fNumValidationEntries += fValidationDatasets[i].GetRows(); - } + for (std::size_t i = 0; i < numSets; i++) { + fNumTrainingEntries += fTrainingRows[i].size(); + fNumValidationEntries += fValidationRows[i].size(); } } diff --git a/tree/ml/src/RDatasetLoader.hxx b/tree/ml/src/RDatasetLoader.hxx index cb04093ba278d..7200557e7987d 100644 --- a/tree/ml/src/RDatasetLoader.hxx +++ b/tree/ml/src/RDatasetLoader.hxx @@ -43,11 +43,11 @@ private: float fVecPadding; std::size_t fNumDatasetCols; - std::vector fTrainingDatasets; - std::vector fValidationDatasets; + RFlat2DMatrix fDataset; - RFlat2DMatrix fTrainingDataset; - RFlat2DMatrix fValidationDataset; + // rows of fDataset in the training and validation sets, one set per dataframe or a single one if concatenated + std::vector> fTrainingRows; + std::vector> fValidationRows; std::size_t fNumTrainingEntries; std::size_t fNumValidationEntries; @@ -65,17 +65,19 @@ public: const std::vector &cols, const std::vector &vecSizes = {}, const float vecPadding = 0.0, bool shuffle = true, const std::size_t setSeed = 0); - RFlat2DMatrix LoadDataframe(ROOT::RDF::RNode &rdf); + std::vector CountEntries(std::vector &rdfs); - void SplitDataset(RFlat2DMatrix &Dataset, RFlat2DMatrix &TrainingDataset, RFlat2DMatrix &ValidationDataset); + void LoadDataframe(ROOT::RDF::RNode &rdf, std::size_t NumEntries); + + void + SplitRows(std::size_t firstRow, std::vector &TrainingRows, std::vector &ValidationRows); void SplitDatasets(bool concatenate); - std::vector ReleaseTrainingDatasets() { return std::move(fTrainingDatasets); } - std::vector ReleaseValidationDatasets() { return std::move(fValidationDatasets); } + RFlat2DMatrix ReleaseDataset() { return std::move(fDataset); } - RFlat2DMatrix ReleaseTrainingDataset() { return std::move(fTrainingDataset); } - RFlat2DMatrix ReleaseValidationDataset() { return std::move(fValidationDataset); } + std::vector> ReleaseTrainingRows() { return std::move(fTrainingRows); } + std::vector> ReleaseValidationRows() { return std::move(fValidationRows); } std::size_t GetNumTrainingEntries() { return fNumTrainingEntries; } std::size_t GetNumValidationEntries() { return fNumValidationEntries; } diff --git a/tree/ml/src/RFlat2DMatrixOperators.cxx b/tree/ml/src/RFlat2DMatrixOperators.cxx index 839e0b0cb2684..0ee6569ca4a98 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.cxx +++ b/tree/ml/src/RFlat2DMatrixOperators.cxx @@ -16,6 +16,22 @@ std::vector RFlat2DMatrixOperators::MakePermutation(std::size_t row return {}; } + // make an identity permutation map + std::vector indices(rows); + std::iota(indices.begin(), indices.end(), 0); + + // shuffle the identity permutation to create a new permutation + ShuffleIndices(indices); + + return indices; +} + +void RFlat2DMatrixOperators::ShuffleIndices(std::vector &indices) +{ + if (!fShuffle) { + return; + } + std::random_device rd; std::mt19937 g; @@ -25,14 +41,7 @@ std::vector RFlat2DMatrixOperators::MakePermutation(std::size_t row g.seed(fSetSeed); } - // make an identity permutation map - std::vector indices(rows); - std::iota(indices.begin(), indices.end(), 0); - - // shuffle the identity permutation to create a new permutation std::shuffle(indices.begin(), indices.end(), g); - - return indices; } void RFlat2DMatrixOperators::AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix &Tensor, diff --git a/tree/ml/src/RFlat2DMatrixOperators.hxx b/tree/ml/src/RFlat2DMatrixOperators.hxx index 12dbd3f65a6ba..52e3ea5cfa79e 100644 --- a/tree/ml/src/RFlat2DMatrixOperators.hxx +++ b/tree/ml/src/RFlat2DMatrixOperators.hxx @@ -39,6 +39,9 @@ public: // returns the order in which to read the rows of a tensor, empty when shuffling is off std::vector MakePermutation(std::size_t rows); + // shuffles the indices in place, nothing when shuffling is off + void ShuffleIndices(std::vector &indices); + // appends the rows [begin, end) of Tensor to Dest, read in the order of the permutation unless it is empty static void AppendRows(RFlat2DMatrix &Dest, const RFlat2DMatrix &Tensor, const std::vector &permutation, std::size_t begin, std::size_t end); diff --git a/tree/ml/src/RSampler.cxx b/tree/ml/src/RSampler.cxx index de004baf763a6..760c7b75a2747 100644 --- a/tree/ml/src/RSampler.cxx +++ b/tree/ml/src/RSampler.cxx @@ -12,9 +12,11 @@ namespace ROOT::Experimental::Internal::ML { RSampler::~RSampler() = default; -RSampler::RSampler(std::vector &datasets, const std::string &sampleType, float sampleRatio, - bool replacement, bool shuffle, std::size_t setSeed) - : fDatasets(datasets), +RSampler::RSampler(const RFlat2DMatrix &dataset, std::vector> &rows, + const std::string &sampleType, float sampleRatio, bool replacement, bool shuffle, + std::size_t setSeed) + : fDataset(dataset), + fRows(rows), fSampleType(sampleType), fSampleRatio(sampleRatio), fReplacement(replacement), @@ -54,7 +56,7 @@ void RSampler::Sampler(RFlat2DMatrix &SampledTensor) /// \brief Calculate fNumEntries and major/minor variables for the random undersampler void RSampler::SetupRandomUndersampler() { - if (fDatasets[0].GetRows() > fDatasets[1].GetRows()) { + if (fRows[0].size() > fRows[1].size()) { fMajor = 0; fMinor = 1; } else { @@ -62,8 +64,8 @@ void RSampler::SetupRandomUndersampler() fMinor = 0; } - fNumMajor = fDatasets[fMajor].GetRows(); - fNumMinor = fDatasets[fMinor].GetRows(); + fNumMajor = fRows[fMajor].size(); + fNumMinor = fRows[fMinor].size(); fNumResampledMajor = static_cast(fNumMinor / fSampleRatio); if (!fReplacement && fNumResampledMajor > fNumMajor) { auto minRatio = std::to_string(std::round(double(fNumMinor) / double(fNumMajor) * 100.0) / 100.0); @@ -80,7 +82,7 @@ void RSampler::SetupRandomUndersampler() /// \brief Calculate fNumEntries and major/minor variables for the random oversampler void RSampler::SetupRandomOversampler() { - if (fDatasets[0].GetRows() > fDatasets[1].GetRows()) { + if (fRows[0].size() > fRows[1].size()) { fMajor = 0; fMinor = 1; } else { @@ -88,8 +90,8 @@ void RSampler::SetupRandomOversampler() fMinor = 0; } - fNumMajor = fDatasets[fMajor].GetRows(); - fNumMinor = fDatasets[fMinor].GetRows(); + fNumMajor = fRows[fMajor].size(); + fNumMinor = fRows[fMinor].size(); fNumResampledMinor = static_cast(fSampleRatio * fNumMajor); fNumEntries = fNumMajor + fNumResampledMinor; } @@ -107,7 +109,7 @@ void RSampler::RandomUndersampler(RFlat2DMatrix &ShuffledTensor) SampleWithoutReplacement(fNumResampledMajor, fNumMajor); } - GatherSamples(ShuffledTensor, fDatasets[fMajor], fNumResampledMajor, fDatasets[fMinor]); + GatherSamples(ShuffledTensor, fRows[fMajor], fNumResampledMajor, fRows[fMinor]); } ////////////////////////////////////////////////////////////////////////// @@ -117,19 +119,19 @@ void RSampler::RandomOversampler(RFlat2DMatrix &ShuffledTensor) { SampleWithReplacement(fNumResampledMinor, fNumMinor); - GatherSamples(ShuffledTensor, fDatasets[fMinor], fNumResampledMinor, fDatasets[fMajor]); + GatherSamples(ShuffledTensor, fRows[fMinor], fNumResampledMinor, fRows[fMajor]); } ////////////////////////////////////////////////////////////////////////// /// \brief Fill the sampled entries: the resampled entries followed by all the other ones, in shuffled order /// \param[in] ShuffledTensor Tensor with all the sampled entries -/// \param[in] ResampledDataset Dataset the entries in fSamples are taken from +/// \param[in] ResampledRows Rows of the class the entries in fSamples are taken from /// \param[in] numResampled Number of resampled entries -/// \param[in] OtherDataset Dataset whose entries are all taken -void RSampler::GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix &ResampledDataset, - std::size_t numResampled, const RFlat2DMatrix &OtherDataset) +/// \param[in] OtherRows Rows of the class whose entries are all taken +void RSampler::GatherSamples(RFlat2DMatrix &ShuffledTensor, const std::vector &ResampledRows, + std::size_t numResampled, const std::vector &OtherRows) { - std::size_t cols = fDatasets[0].GetCols(); + std::size_t cols = fDataset.GetCols(); const auto permutation = fTensorOperators->MakePermutation(fNumEntries); // read the entries straight from the datasets, instead of concatenating and shuffling copies of them @@ -137,9 +139,8 @@ void RSampler::GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix float *out = ShuffledTensor.GetData(); for (std::size_t i = 0; i < fNumEntries; i++) { const std::size_t entry = permutation.empty() ? i : permutation[i]; - const float *row = entry < numResampled ? ResampledDataset.GetData() + fSamples[entry] * cols - : OtherDataset.GetData() + (entry - numResampled) * cols; - out = std::copy_n(row, cols, out); + const std::size_t row = entry < numResampled ? ResampledRows[fSamples[entry]] : OtherRows[entry - numResampled]; + out = std::copy_n(fDataset.GetData() + row * cols, cols, out); } } diff --git a/tree/ml/src/RSampler.hxx b/tree/ml/src/RSampler.hxx index 7f35298df31ea..ed0b573663a70 100644 --- a/tree/ml/src/RSampler.hxx +++ b/tree/ml/src/RSampler.hxx @@ -31,7 +31,9 @@ namespace ROOT::Experimental::Internal::ML { class RSampler { private: - std::vector &fDatasets; + const RFlat2DMatrix &fDataset; + // rows of fDataset in each class + std::vector> &fRows; std::string fSampleType; float fSampleRatio; bool fReplacement; @@ -51,8 +53,8 @@ private: std::unique_ptr fTensorOperators; public: - RSampler(std::vector &datasets, const std::string &sampleType, float sampleRatio, - bool replacement = false, bool shuffle = true, std::size_t setSeed = 0); + RSampler(const RFlat2DMatrix &dataset, std::vector> &rows, const std::string &sampleType, + float sampleRatio, bool replacement = false, bool shuffle = true, std::size_t setSeed = 0); ~RSampler(); @@ -68,8 +70,8 @@ public: void RandomOversampler(RFlat2DMatrix &ShuffledTensor); - void GatherSamples(RFlat2DMatrix &ShuffledTensor, const RFlat2DMatrix &ResampledDataset, std::size_t numResampled, - const RFlat2DMatrix &OtherDataset); + void GatherSamples(RFlat2DMatrix &ShuffledTensor, const std::vector &ResampledRows, + std::size_t numResampled, const std::vector &OtherRows); void SampleWithReplacement(std::size_t n_samples, std::size_t max); From f137c16d08650958ac41985f73ba2c4b156b62c5 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Mon, 19 Oct 2026 12:06:10 +0200 Subject: [PATCH 8/8] [RDF] Opt.: Reuse the value addresses of the custom loader action for every entry `RActionCustomLoader::CallExec` allocated a new vector of value addresses for every entry. Assisted-by: Claude Opus 5.5 --- tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx b/tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx index eb79e370bca1c..a38564319aa69 100644 --- a/tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx +++ b/tree/dataframe/inc/ROOT/RDF/RActionCustomLoader.hxx @@ -34,6 +34,9 @@ class R__CLING_PTRCHECK(off) RActionCustomLoader final : public RActionBase { /// Column readers per slot and per input column std::vector> fValues; + /// Value addresses of the current entry per slot, reused for every entry + std::vector> fValueAddresses; + // Types of input columns std::vector fColTypeIDs; @@ -47,6 +50,7 @@ public: fHelper(std::move(h)), fPrevNodes{std::static_pointer_cast(pd)}, fValues(GetNSlots()), + fValueAddresses(GetNSlots()), fColTypeIDs(colTypeIDs) { fLoopManager->Register(this); @@ -93,11 +97,11 @@ public: void CallExec(unsigned int slot, Long64_t entry) { - std::vector untypedValues; + auto &untypedValues = fValueAddresses[slot]; auto nReaders = fValues[slot].size(); - untypedValues.reserve(nReaders); + untypedValues.resize(nReaders); for (decltype(nReaders) readerIdx{}; readerIdx < nReaders; readerIdx++) - untypedValues.push_back(GetValue(slot, readerIdx, entry)); + untypedValues[readerIdx] = GetValue(slot, readerIdx, entry); fHelper.Exec(slot, untypedValues); }