EIC code displayed by LXR master/​include/​root/​TMVA/​RBatchLoader.hxx	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-01-18 10:11:04

 #ifndef TMVA_RBatchLoader
 #define TMVA_RBatchLoader
 
 #include <iostream>
 #include <vector>
 #include <memory>
 
 // Imports for threading
 #include <queue>
 #include <mutex>
 #include <condition_variable>
 
 #include "TMVA/RTensor.hxx"
 #include "TMVA/Tools.h"
 #include "TRandom3.h"
 
 namespace TMVA {
 namespace Experimental {
 namespace Internal {
 
 class RBatchLoader {
 private:
    std::size_t fBatchSize;
    std::size_t fNumColumns;
    std::size_t fMaxBatches;
 
    bool fIsActive = false;
    TMVA::RandomGenerator<TRandom3> fRng = TMVA::RandomGenerator<TRandom3>(0);
 
    std::mutex fBatchLock;
    std::condition_variable fBatchCondition;
 
    std::queue<std::unique_ptr<TMVA::Experimental::RTensor<float>>> fTrainingBatchQueue;
    std::vector<std::unique_ptr<TMVA::Experimental::RTensor<float>>> fValidationBatches;
    std::unique_ptr<TMVA::Experimental::RTensor<float>> fCurrentBatch;
 
    std::size_t fValidationIdx = 0;
 
    TMVA::Experimental::RTensor<float> fEmptyTensor = TMVA::Experimental::RTensor<float>({0});
 
 public:
    RBatchLoader(const std::size_t batchSize, const std::size_t numColumns, const std::size_t maxBatches)
       : fBatchSize(batchSize), fNumColumns(numColumns), fMaxBatches(maxBatches)
    {
    }
 
    ~RBatchLoader() { DeActivate(); }
 
 public:
    /// \brief Return a batch of data as a unique pointer.
    /// After the batch has been processed, it should be distroyed.
    /// \return Training batch
    const TMVA::Experimental::RTensor<float> &GetTrainBatch()
    {
       std::unique_lock<std::mutex> lock(fBatchLock);
       fBatchCondition.wait(lock, [this]() { return !fTrainingBatchQueue.empty() || !fIsActive; });
 
       if (fTrainingBatchQueue.empty()) {
          fCurrentBatch = std::make_unique<TMVA::Experimental::RTensor<float>>(std::vector<std::size_t>({0}));
          return *fCurrentBatch;
       }
 
       fCurrentBatch = std::move(fTrainingBatchQueue.front());
       fTrainingBatchQueue.pop();
 
       fBatchCondition.notify_all();
 
       return *fCurrentBatch;
    }
 
    /// \brief Returns a batch of data for validation
    /// The owner of this batch has to be with the RBatchLoader.
    /// This is because the same validation batches should be used in all epochs.
    /// \return Validation batch
    const TMVA::Experimental::RTensor<float> &GetValidationBatch()
    {
       if (HasValidationData()) {
          return *fValidationBatches[fValidationIdx++].get();
       }
 
       return fEmptyTensor;
    }
 
    /// \brief Checks if there are more training batches available
    /// \return
    bool HasTrainData()
    {
       {
          std::unique_lock<std::mutex> lock(fBatchLock);
          if (!fTrainingBatchQueue.empty() || fIsActive)
             return true;
       }
 
       return false;
    }
 
    /// \brief Checks if there are more training batches available
    /// \return
    bool HasValidationData()
    {
       std::unique_lock<std::mutex> lock(fBatchLock);
       return fValidationIdx < fValidationBatches.size();
    }
 
    /// \brief Activate the batchloader so it will accept chunks to batch
    void Activate()
    {
       {
          std::lock_guard<std::mutex> lock(fBatchLock);
          fIsActive = true;
       }
       fBatchCondition.notify_all();
    }
 
    /// \brief DeActivate the batchloader. This means that no more batches are created.
    /// Batches can still be returned if they are already loaded
    void DeActivate()
    {
       {
          std::lock_guard<std::mutex> lock(fBatchLock);
          fIsActive = false;
       }
       fBatchCondition.notify_all();
    }
 
    /// \brief Create a batch filled with the events on the given idx
    /// \param chunkTensor
    /// \param idx
    /// \return
    std::unique_ptr<TMVA::Experimental::RTensor<float>>
    CreateBatch(const TMVA::Experimental::RTensor<float> &chunkTensor, const std::vector<std::size_t> idx)
    {
       auto batch =
          std::make_unique<TMVA::Experimental::RTensor<float>>(std::vector<std::size_t>({fBatchSize, fNumColumns}));
 
       for (std::size_t i = 0; i < fBatchSize; i++) {
          std::copy(chunkTensor.GetData() + (idx[i] * fNumColumns), chunkTensor.GetData() + ((idx[i] + 1) * fNumColumns),
                    batch->GetData() + i * fNumColumns);
       }
 
       return batch;
    }
 
    /// \brief Create training batches from the given chunk of data based on the given event indices
    /// Batches are added to the training queue of batches
    /// The eventIndices can be shuffled to ensure random order for each epoch
    /// \param chunkTensor
    /// \param eventIndices
    /// \param shuffle
    void CreateTrainingBatches(const TMVA::Experimental::RTensor<float> &chunkTensor,
                               std::vector<std::size_t> eventIndices, const bool shuffle = true)
    {
       // Wait until less than a full chunk of batches are in the queue before loading splitting the next chunk into
       // batches
       {
          std::unique_lock<std::mutex> lock(fBatchLock);
          fBatchCondition.wait(lock, [this]() { return (fTrainingBatchQueue.size() < fMaxBatches) || !fIsActive; });
          if (!fIsActive)
             return;
       }
 
       if (shuffle)
          std::shuffle(eventIndices.begin(), eventIndices.end(), fRng); // Shuffle the order of idx
 
       std::vector<std::unique_ptr<TMVA::Experimental::RTensor<float>>> batches;
 
       // Create tasks of fBatchSize untill all idx are used
       for (std::size_t start = 0; (start + fBatchSize) <= eventIndices.size(); start += fBatchSize) {
 
          // Grab the first fBatchSize indices from the
          std::vector<std::size_t> idx;
          for (std::size_t i = start; i < (start + fBatchSize); i++) {
             idx.push_back(eventIndices[i]);
          }
 
          // Fill a batch
          batches.emplace_back(CreateBatch(chunkTensor, idx));
       }
 
       {
          std::unique_lock<std::mutex> lock(fBatchLock);
          for (std::size_t i = 0; i < batches.size(); i++) {
             fTrainingBatchQueue.push(std::move(batches[i]));
          }
       }
 
       fBatchCondition.notify_one();
    }
 
    /// \brief Create validation batches from the given chunk based on the given event indices
    /// Batches are added to the vector of validation batches
    /// \param chunkTensor
    /// \param eventIndices
    void CreateValidationBatches(const TMVA::Experimental::RTensor<float> &chunkTensor,
                                 const std::vector<std::size_t> eventIndices)
    {
       // Create tasks of fBatchSize untill all idx are used
       for (std::size_t start = 0; (start + fBatchSize) <= eventIndices.size(); start += fBatchSize) {
 
          std::vector<std::size_t> idx;
 
          for (std::size_t i = start; i < (start + fBatchSize); i++) {
             idx.push_back(eventIndices[i]);
          }
 
          {
             std::unique_lock<std::mutex> lock(fBatchLock);
             fValidationBatches.emplace_back(CreateBatch(chunkTensor, idx));
          }
       }
    }
 
    /// \brief Reset the validation process
    void StartValidation()
    {
       std::unique_lock<std::mutex> lock(fBatchLock);
       fValidationIdx = 0;
    }
 };
 
 } // namespace Internal
 } // namespace Experimental
 } // namespace TMVA
 
 #endif // TMVA_RBatchLoader

This page was automatically generated by the 2.3.7 LXR engine. The LXR team