batch_training.h

Go to the documentation of this file.

#ifndef META_BATCH_TRAINING_H_
#define META_BATCH_TRAINING_H_

#include <algorithm>
#include <random>
#include <vector>

#include "logging/logger.h"
#include "meta.h"

namespace meta
{
namespace classify
{

template <class Index, class Classifier>
void batch_train(Index& idx, Classifier& cls,
                 const std::vector<doc_id>& training_set, uint64_t batch_size)
{
    auto docs = training_set;
    std::mt19937 gen(std::random_device{}());
    std::shuffle(docs.begin(), docs.end(), gen);

    // integer-math ceil(docs.size() / batch_size)
    auto num_batches = (docs.size() + batch_size - 1) / batch_size;
    for (uint64_t i = 0; i < num_batches; ++i)
    {
        LOG(progress) << "\rTraining batch " << i + 1 << "/" << num_batches
                      << ENDLG;
        auto end = std::min<uint64_t>((i + 1) * batch_size, docs.size());
        std::vector<doc_id> batch{docs.begin() + (i * batch_size),
                                  docs.begin() + end};
        idx.clear_cache();
        cls.train(batch);
    }
    LOG(progress) << '\n' << ENDLG;
}

}
}
#endif