ROOT/ML/RBatchLoader.hxx

0001 // Author: Dante Niewenhuis, VU Amsterdam 07/2023
0002 // Author: Kristupas Pranckietis, Vilnius University 05/2024
0003 // Author: Nopphakorn Subsa-Ard, King Mongkut's University of Technology Thonburi (KMUTT) (TH) 08/2024
0004 // Author: Vincenzo Eduardo Padulano, CERN 10/2024
0005 // Author: Martin Føll, University of Oslo (UiO) & CERN 05/2025
0006 // Author: Silia Taider, CERN 02/2026
0007
0008 /*************************************************************************
0009  * Copyright (C) 1995-2025, Rene Brun and Fons Rademakers.               *
0010  * All rights reserved.                                                  *
0011  *                                                                       *
0012  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0013  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0014  *************************************************************************/
0015
0016 #ifndef ROOT_INTERNAL_ML_RBATCHLOADER
0017 #define ROOT_INTERNAL_ML_RBATCHLOADER
0018
0019 #include <condition_variable>
0020 #include <memory>
0021 #include <mutex>
0022 #include <queue>
0023 #include <string>
0024 #include <vector>
0025
0026 #include "ROOT/ML/RFlat2DMatrix.hxx"
0027
0028 namespace ROOT::Experimental::Internal::ML {
0029 /**
0030  \class ROOT::Experimental::Internal::ML::RBatchLoader
0031
0032 \brief Building and loading the batches from loaded chunks in RChunkLoader
0033
0034 In this class the chunks that are loaded into memory (see RChunkLoader) are split into batches used in the ML training
0035 which are loaded into a queue. This is done for both the training and validation chunks separately.
0036 */
0037
0038 class RBatchLoader {
0039 private:
0040    std::size_t fBatchSize;
0041    // needed for calculating the total number of batch columns when vectors columns are present
0042    std::vector<std::string> fCols;
0043    std::mutex &fLock;
0044    std::condition_variable &fCV;
0045    std::vector<std::size_t> fVecSizes;
0046    std::size_t fSumVecSizes;
0047    std::size_t fNumColumns;
0048    std::size_t fNumEntries;
0049    bool fDropRemainder;
0050
0051    std::size_t fNumFullBatches;
0052    std::size_t fNumBatches;
0053    std::size_t fLeftoverBatchSize;
0054
0055    bool fIsActive = false;
0056    bool fProducerDone = true;
0057
0058    // queues of flattened tensors (rows * cols)
0059    std::queue<std::unique_ptr<RFlat2DMatrix>> fBatchQueue;
0060
0061    // current batch that is loaded into memory
0062    std::unique_ptr<RFlat2DMatrix> fCurrentBatch;
0063
0064    // primary and secondary leftover batches used to create batches from a chunk
0065    std::unique_ptr<RFlat2DMatrix> fPrimaryLeftoverBatch;
0066    std::unique_ptr<RFlat2DMatrix> fSecondaryLeftoverBatch;
0067
0068 public:
0069    RBatchLoader(std::size_t batchSize, const std::vector<std::string> &cols, std::mutex &sharedMutex,
0070                 std::condition_variable &sharedCV, const std::vector<std::size_t> &vecSizes = {},
0071                 std::size_t numEntries = 0, bool dropRemainder = false);
0072
0073    void Activate();
0074    void DeActivate();
0075    void Reset();
0076    void MarkProducerDone();
0077
0078    std::unique_ptr<RFlat2DMatrix> CreateBatch(RFlat2DMatrix &chunkTensor, std::size_t idxs);
0079    RFlat2DMatrix GetBatch();
0080    void CreateBatches(RFlat2DMatrix &chunkTensor, bool isLastBatch);
0081    void RecalculateBatchCounts(std::size_t numEntries);
0082
0083    bool isProducerDone() { return fProducerDone; }
0084    std::size_t GetNumBatches() { return fNumBatches; }
0085    std::size_t GetNumEntries() { return fNumEntries; }
0086    std::size_t GetNumRemainderRows() { return fLeftoverBatchSize; }
0087    std::size_t GetNumBatchQueue() { return fBatchQueue.size(); }
0088 };
0089
0090 } // namespace ROOT::Experimental::Internal::ML
0091
0092 #endif // ROOT_INTERNAL_ML_RBATCHLOADER