BmnGemStripMLPModel.h

Go to the documentation of this file.

#ifndef BMNGEMSTRIPMLPMODEL_H
#define BMNGEMSTRIPMLPMODEL_H
 
// Neural network inference for the GEM strip simulation (no external libraries).
//
// BmnGemStripMLPModel - residual MLP trained in Keras (build_generator_mlp):
//
//   input(input_dim) -> Dense(width) -> LeakyReLU
//                    -> [ LayerNorm -> Dense(width) -> LeakyReLU -> Dense(width) -> Add ] x n_blocks
//                    -> LayerNorm -> Dense(output_dim) -> sigmoid (or linear)
//
// width and n_blocks are read from the model file, so a retrained model with another
// width / n_blocks does not need recompilation.
//
// Model file: binary .bmnmlp written by Converter_KerasToBmnMLP.py
// (weights, architecture parameters, feature names and test samples for the self-test on loading).
//
// Usage:
//   auto model = BmnGemStripMLPModel::GetShared();              // default path, shared between modules
//   auto model = BmnGemStripMLPModel::GetShared("my.bmnmlp");   // explicit path
//   BmnGemStripMLPModel::SetDefaultPath("$VMCWORKDIR/parameters/gem/ANNModels/other.bmnmlp");
//   model->Predict(input, output);                               // input_dim floats -> output_dim floats
//   std::vector<float> out = model->Predict(input_vector);      // the same, empty vector if failed
//
// The interface is the same as in BmnGemStripFdeepModel (frugally-deep models).
//
// After loading the model is not modified: Predict() can be called from several threads at once.
 
#include <iostream>
#include <memory>
#include <string>
#include <vector>
 
class BmnGemStripMLPModel
{
 
  public:
    // Limits (stack buffers in the inference) -----------------------------------
    static constexpr int kMaxWidth = 1024;   // max hidden width
    static constexpr int kMaxDim = 1024;     // max input / output size
 
    // Output activation of the last Dense layer
    enum OutputActivation
    {
        kLinearOutput = 0,
        kSigmoidOutput = 1
    };
 
    // Constructors -------------------------------------------------------------
    BmnGemStripMLPModel();
    explicit BmnGemStripMLPModel(const std::string& path, bool verbose = true);   // Load(path)
    //--------------------------------------------------------------------------
 
    // Destructor
    virtual ~BmnGemStripMLPModel();
 
    // Loading ------------------------------------------------------------------
    // Reads the model file, checks it (format, sizes, checksum) and runs the self-test.
    // Returns false (and the model stays empty) if anything is wrong; the reason is printed (LOG(error)) if verbose.
    bool Load(const std::string& path, bool verbose = true);
    void Clear();
    bool IsLoaded() const { return fLoaded; }
    //--------------------------------------------------------------------------
 
    // Model parameters ---------------------------------------------------------
    int GetInputDim() const { return fInputDim; }
    int GetWidth() const { return fWidth; }
    int GetNBlocks() const { return fNBlocks; }
    int GetOutputDim() const { return fOutputDim; }
    int GetOutputActivation() const { return fOutputActivation; }
    float GetLayerNormEpsilon() const { return fLayerNormEps; }
    float GetLeakyReLUSlope() const { return fLeakyReLUSlope; }
    size_t GetNParams() const { return fParams.size(); }
 
    const std::vector<std::string>& GetFeatureNames() const { return fFeatureNames; }
    int GetFeatureIndex(const std::string& name) const;   // -1 if there is no such feature
    // true if the feature names are exactly 'expected' (the same order); the reason is written to 'why'
    bool CheckFeatureNames(const std::vector<std::string>& expected, std::string* why = nullptr) const;
 
    const std::string& GetFilePath() const { return fFilePath; }
    const std::string& GetArchHash() const { return fArchHash; }
    const std::string& GetSourceName() const { return fSourceName; }   // name of the original .keras file
    const std::string& GetCreated() const { return fCreated; }         // conversion date
 
    void Print(std::ostream& os = std::cout) const;   // short summary
    //--------------------------------------------------------------------------
 
    // Inference ----------------------------------------------------------------
    // One sample: input[GetInputDim()] -> output[GetOutputDim()]
    bool Predict(const float* input, float* output) const;
    // The same with vectors: input.size() must be GetInputDim(); returns an empty vector if failed
    std::vector<float> Predict(const std::vector<float>& input) const;
 
    // Several samples with a common first part of the input (the first Dense is computed for it only once):
    //   input of row r = [ common[0 .. n_common-1], specific[r*n_spec .. r*n_spec + n_spec-1] ],
    //   n_spec = GetInputDim() - n_common,
    //   output[r*GetOutputDim() .. ] - output of row r.
    bool PredictSplit(const float* common, int n_common, const float* specific, int n_rows, float* output) const;
    //--------------------------------------------------------------------------
 
    // Self-test: test samples from the model file (Keras outputs) ---------------
    int GetNTests() const { return fNTests; }
    float GetTestTolerance() const { return fTestTolerance; }
    bool RunSelfTest(float* max_abs_diff = nullptr, bool verbose = false) const;   // verbose: result by LOG(info)
    const float* GetTestInput(int itest) const;                                    // nullptr if out of range
    const float* GetTestOutput(int itest) const;                                   // nullptr if out of range
    //--------------------------------------------------------------------------
 
    // Shared models (one copy per file in the process) --------------------------
    // path: "" - the default path; environment variables ($VMCWORKDIR, ~, ...) are expanded (gSystem->ExpandPathName).
    // The same file given by different strings is loaded once (the cache key is the expanded path).
    // A file that failed to load is not reloaded on each call (the warning is printed once);
    // ClearSharedCache() resets this.
    static std::shared_ptr<const BmnGemStripMLPModel> GetShared(const std::string& path = "");
    static void SetDefaultPath(const std::string& path);
    static std::string GetDefaultPath();
    static std::string GetBuiltinDefaultPath();   // "$VMCWORKDIR/parameters/gem/ANNModels/GEM_DefaultModel.bmnmlp"
    static void ClearSharedCache();
    //--------------------------------------------------------------------------
 
  private:
    // h: output of the first Dense (bias included, before the activation) -> output
    void ForwardFromFirstDense(float* h, float* output) const;
 
    bool Parse(const std::vector<unsigned char>& buffer, std::string& error);
 
  private:
    bool fLoaded;
 
    int fInputDim;
    int fWidth;
    int fNBlocks;
    int fOutputDim;
    int fOutputActivation;
    float fLayerNormEps;
    float fLeakyReLUSlope;
 
    std::string fFilePath;
    std::string fArchHash;
    std::string fSourceName;
    std::string fCreated;
    std::vector<std::string> fFeatureNames;
 
    // All weights in one array (the order of the file), positions of the parts:
    std::vector<float> fParams;
    size_t fOffDenseInW;    // [input_dim][width]
    size_t fOffDenseInB;    // [width]
    size_t fOffBlocks;      // n_blocks x { ln_gamma[w], ln_beta[w], W1[w][w], b1[w], W2[w][w], b2[w] }
    size_t fBlockSize;      // floats in one block
    size_t fOffFinalLnG;    // [width]
    size_t fOffFinalLnB;    // [width]
    size_t fOffDenseOutW;   // [width][output_dim]
    size_t fOffDenseOutB;   // [output_dim]
 
    // Test samples
    int fNTests;
    float fTestTolerance;
    std::vector<float> fTestInputs;    // [n_tests][input_dim]
    std::vector<float> fTestOutputs;   // [n_tests][output_dim]
};
//------------------------------------------------------------------------------
 
#endif /* BMNGEMSTRIPMLPMODEL_H */