BmnGemStripFdeepModel.cxx

Go to the documentation of this file.

#include "BmnGemStripFdeepModel.h"
 
#include "FairLogger.h"
#include "TString.h"
#include "TSystem.h"
 
#include <algorithm>
#include <map>
#include <mutex>
#include <set>
 
// frugally-deep (optional) ------------------------------------------------------
// BMN_WITH_FDEEP = 1 : fdeep models can be loaded (fdeep found by include)
// BMN_WITH_FDEEP = 0 : built without fdeep, no model can be loaded (the ANN methods use the standard method)
// Can be forced from the build: -DBMN_WITH_FDEEP=0 (or =1)
#ifndef BMN_WITH_FDEEP
#if __has_include("fdeep/fdeep.hpp")
#define BMN_WITH_FDEEP 1
#else
#define BMN_WITH_FDEEP 0
#endif
#endif
 
#if BMN_WITH_FDEEP
#define FDEEP_FLOAT_TYPE float
#include "fdeep/fdeep.hpp"
#endif
//------------------------------------------------------------------------------
 
namespace
{
const char* const kBuiltinDefaultPath = "$VMCWORKDIR/parameters/gem/ANNModels/GEM_DefaultModel.json";
 
// Path with environment variables ($VMCWORKDIR, ~, ...) -> real path (ROOT: gSystem->ExpandPathName).
// ok = false if it cannot be expanded (an environment variable is not set); the path is returned as given then.
std::string ExpandPathName(const std::string& path, bool* ok = nullptr)
{
    TString expanded(path.c_str());
    const bool failed = gSystem->ExpandPathName(expanded);   // kTRUE in case of error
    if (ok)
        *ok = !failed;
    return failed ? path : std::string(expanded.Data());
}
 
// Shared models ---------------------------------------------------------------
struct SharedState
{
    std::mutex mutex;
    std::string default_path = kBuiltinDefaultPath;
    std::map<std::string, std::shared_ptr<const BmnGemStripFdeepModel>> models;   // key - expanded path
    std::set<std::string> failed;                                                 // paths that failed to load
    std::map<std::string, std::string> expanded;                                  // path as given -> expanded path
};
 
SharedState& Shared()
{
    static SharedState state;
    return state;
}
//==============================================================================
}   // namespace
//------------------------------------------------------------------------------
 
// fdeep model and the shape of its input (only here: the header does not need fdeep)
struct BmnGemStripFdeepModel::Impl
{
#if BMN_WITH_FDEEP
    Impl(fdeep::model&& m, const fdeep::tensor_shape& shape)
        : model(std::move(m))
        , input_shape(shape)
    {}
    fdeep::model model;
    fdeep::tensor_shape input_shape;
#endif
};
//------------------------------------------------------------------------------
 
BmnGemStripFdeepModel::BmnGemStripFdeepModel()
    : fLoaded(false)
    , fInputDim(0)
    , fOutputDim(0)
{}
 
BmnGemStripFdeepModel::BmnGemStripFdeepModel(const std::string& path, bool verbose)
    : BmnGemStripFdeepModel()
{
    Load(path, verbose);
}
 
BmnGemStripFdeepModel::~BmnGemStripFdeepModel() {}
 
void BmnGemStripFdeepModel::Clear()
{
    fImpl.reset();
    fLoaded = false;
    fInputDim = 0;
    fOutputDim = 0;
    fFilePath.clear();
}
 
bool BmnGemStripFdeepModel::IsBuiltWithFdeep()
{
    return BMN_WITH_FDEEP != 0;
}
//------------------------------------------------------------------------------
 
bool BmnGemStripFdeepModel::Load(const std::string& path, bool verbose)
{
    Clear();
 
    bool expanded_ok = true;
    const std::string full_path = ExpandPathName(path, &expanded_ok);
 
    auto fail = [&](const std::string& reason) {
        if (verbose)
            LOG(error) << "BmnGemStripFdeepModel: cannot load model " << full_path << ": " << reason;
        Clear();
        return false;
    };
 
#if BMN_WITH_FDEEP
    if (gSystem->AccessPathName(full_path.c_str(), kFileExists)) {
        std::string reason = "file not found";
        if (!expanded_ok)
            reason += " (an environment variable in the path is not set)";
        return fail(reason);
    }
 
    try {
        // verify = false: the test data of the file are not checked (fast loading)
        fdeep::model model = fdeep::load_model(full_path, false, fdeep::dev_null_logger);
 
        if (model.get_input_shapes().size() != 1 || model.get_output_shapes().size() != 1)
            return fail("the model must have one input and one output");
 
        // Sizes of the input and the output: one prediction on zero input
        const fdeep::tensors dummy = model.generate_dummy_inputs();
        const fdeep::tensors out = model.predict(dummy);
        if (dummy.size() != 1 || out.size() != 1)
            return fail("unexpected number of input / output tensors");
 
        const fdeep::tensor_shape input_shape = dummy.front().shape();
        fInputDim = static_cast<int>(input_shape.volume());
        fOutputDim = static_cast<int>(out.front().as_vector()->size());
        fImpl.reset(new Impl(std::move(model), input_shape));
    } catch (const std::exception& e) {
        return fail(e.what());
    }
 
    fFilePath = full_path;
    fLoaded = true;
    return true;
#else
    (void)expanded_ok;
    return fail("the library is built without frugally-deep");
#endif
}
//------------------------------------------------------------------------------
 
bool BmnGemStripFdeepModel::Predict(const float* input, float* output) const
{
    if (!fLoaded || !fImpl || !input || !output)
        return false;
 
#if BMN_WITH_FDEEP
    try {
        fdeep::tensors in;
        in.emplace_back(fImpl->input_shape, std::vector<float>(input, input + fInputDim));
        const fdeep::tensors out = fImpl->model.predict(in);
        if (out.empty())
            return false;
        const auto& result = *out.front().as_vector();   // fdeep vector may use its own allocator
        if (static_cast<int>(result.size()) != fOutputDim)
            return false;
        std::copy(result.begin(), result.end(), output);
        return true;
    } catch (const std::exception& e) {
        LOG(error) << "BmnGemStripFdeepModel: prediction failed (" << fFilePath << "): " << e.what();
        return false;
    }
#else
    return false;
#endif
}
 
std::vector<float> BmnGemStripFdeepModel::Predict(const std::vector<float>& input) const
{
    if (!fLoaded || static_cast<int>(input.size()) != fInputDim)
        return {};
    std::vector<float> output(fOutputDim);
    if (!Predict(input.data(), output.data()))
        return {};
    return output;
}
//------------------------------------------------------------------------------
 
bool BmnGemStripFdeepModel::PredictSplit(const float* common,
                                         int n_common,
                                         const float* specific,
                                         int n_rows,
                                         float* output) const
{
    if (!fLoaded || !output || n_rows < 0 || n_common < 0 || n_common > fInputDim)
        return false;
    const int n_spec = fInputDim - n_common;
    if ((n_common > 0 && !common) || (n_spec > 0 && n_rows > 0 && !specific))
        return false;
 
    // Each row separately: [common, specific of the row]
    std::vector<float> input(fInputDim);
    if (n_common > 0)
        std::copy(common, common + n_common, input.begin());
    for (int r = 0; r < n_rows; ++r) {
        if (n_spec > 0)
            std::copy(specific + static_cast<size_t>(r) * n_spec, specific + static_cast<size_t>(r + 1) * n_spec,
                      input.begin() + n_common);
        if (!Predict(input.data(), output + static_cast<size_t>(r) * fOutputDim))
            return false;
    }
    return true;
}
//------------------------------------------------------------------------------
 
void BmnGemStripFdeepModel::Print(std::ostream& os) const
{
    if (!fLoaded) {
        os << "BmnGemStripFdeepModel: not loaded" << (IsBuiltWithFdeep() ? "" : " (built without frugally-deep)")
           << "\n";
        return;
    }
    os << "BmnGemStripFdeepModel: " << fFilePath << "\n"
       << "  input " << fInputDim << " -> output " << fOutputDim << "\n";
}
//------------------------------------------------------------------------------
 
std::shared_ptr<const BmnGemStripFdeepModel> BmnGemStripFdeepModel::GetShared(const std::string& path)
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
 
    // Cache key - the expanded path (the same file given by different strings is loaded once);
    // the expansion is done once for each path string (GetShared() can be called for each point)
    const std::string& given = path.empty() ? state.default_path : path;
    auto iexp = state.expanded.find(given);
    if (iexp == state.expanded.end())
        iexp = state.expanded.emplace(given, ExpandPathName(given)).first;
    const std::string key = iexp->second;
 
    auto it = state.models.find(key);
    if (it != state.models.end())
        return it->second;
    if (state.failed.count(key))
        return nullptr;
 
    auto model = std::make_shared<BmnGemStripFdeepModel>();
    if (!model->Load(key, true)) {
        LOG(warn) << "BmnGemStripFdeepModel: the model " << key
                  << " is not available (this warning is printed once for the path)";
        state.failed.insert(key);
        return nullptr;
    }
    std::shared_ptr<const BmnGemStripFdeepModel> shared = model;
    state.models[key] = shared;
    return shared;
}
 
void BmnGemStripFdeepModel::SetDefaultPath(const std::string& path)
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    state.default_path = path.empty() ? std::string(kBuiltinDefaultPath) : path;
}
 
std::string BmnGemStripFdeepModel::GetDefaultPath()
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    return state.default_path;
}
 
std::string BmnGemStripFdeepModel::GetBuiltinDefaultPath()
{
    return kBuiltinDefaultPath;
}
 
void BmnGemStripFdeepModel::ClearSharedCache()
{
    SharedState& state = Shared();
    std::lock_guard<std::mutex> lock(state.mutex);
    state.models.clear();   // models in use stay alive while someone holds them
    state.failed.clear();
    state.expanded.clear();
}