root/TMVA/RSofieReader.hxx

0001 /**********************************************************************************
0002  * Project: ROOT - a Root-integrated toolkit for multivariate data analysis       *
0003  * Package: TMVA                                                                  *                                        *
0004  *                                                                                *
0005  * Description:                                                                   *
0006  *                                                                                *
0007  * Authors:                                                                       *
0008  *      Lorenzo Moneta                                  *
0009  *                                                                                *
0010  * Copyright (c) 2022:                                                            *
0011  *      CERN, Switzerland                                                         *
0012  *                                                                                *
0013  **********************************************************************************/
0014
0015
0016 #ifndef TMVA_RSOFIEREADER
0017 #define TMVA_RSOFIEREADER
0018
0019
0020 #include <string>
0021 #include <vector>
0022 #include <memory> // std::unique_ptr
0023 #include <sstream> // std::stringstream
0024 #include <iostream>
0025 #include "TROOT.h"
0026 #include "TSystem.h"
0027 #include "TError.h"
0028 #include "TInterpreter.h"
0029 #include "TUUID.h"
0030 #include "TMVA/RTensor.hxx"
0031 #include "Math/Util.h"
0032
0033 namespace TMVA {
0034 namespace Experimental {
0035
0036
0037
0038
0039 /// TMVA::RSofieReader class for reading external Machine Learning models
0040 /// in ONNX files, Keras .h5  or .keras files or PyTorch .pt files
0041 /// and performing the inference using SOFIE
0042 /// It is reccomended to use ONNX if possible since there is a larger support for
0043 /// model operators.
0044
0045 class RSofieReader  {
0046
0047
0048 public:
0049    /// Dummy constructor which needs model loading  afterwards
0050    RSofieReader() {}
0051    /// Create TMVA model from ONNX file
0052    /// print level can be 0 (minimal) 1 with info , 2 with all ONNX parsing info
0053    RSofieReader(const std::string &path, std::vector<std::vector<size_t>> inputShapes = {}, int verbose = 0)
0054    {
0055       Load(path, inputShapes, verbose);
0056    }
0057
0058    void Load(const std::string &path, std::vector<std::vector<size_t>> inputShapes = {}, int verbose = 0)
0059    {
0060
0061       enum EModelType {kONNX, kKeras, kPt, kROOT, kNotDef}; // type of model
0062       EModelType type = kNotDef;
0063
0064       size_t pos2 = std::string::npos;
0065       if ( (pos2 = path.find(".onnx")) != std::string::npos) {
0066          if (verbose) std::cout << "input model type is ONNX" << std::endl;
0067          type = kONNX;
0068       } else if ( (pos2 = path.find(".h5")) != std::string::npos || (pos2 = path.find(".keras")) != std::string::npos) {
0069          if (verbose) std::cout << "input model type is Keras" << std::endl;
0070          type = kKeras;
0071       } else if ( (pos2 = path.find(".pt")) != std::string::npos) {
0072          if (verbose) std::cout << "input model type is PyTorch" << std::endl;
0073          type = kPt;
0074       } else if ( (pos2 = path.find(".root")) != std::string::npos) {
0075          if (verbose) std::cout << "input model type is ROOT" << std::endl;
0076          type = kROOT;
0077       }
0078
0079       if (type == kNotDef) {
0080          throw std::runtime_error("Input file is not an ONNX or Keras or PyTorch file");
0081       }
0082       auto pos1 = path.rfind("/");
0083       if (pos1 == std::string::npos)
0084          pos1 = 0;
0085       else
0086          pos1 += 1;
0087       std::string modelName = path.substr(pos1,pos2-pos1);
0088       std::string fileType = path.substr(pos2+1, path.length()-pos2-1);
0089       if (verbose) std::cout << "Parsing SOFIE model " << modelName << " of type " << fileType << std::endl;
0090
0091       // append a suffix to headerfile
0092       std::string modelHeader = modelName + "_fromRSofieR.hxx";
0093       std::string modelWeights = modelName + "_fromRSofieR.dat";
0094
0095       // create code for parsing model and generate C++ code for inference
0096       // make it in a separate scope to avoid polluting global interpreter space
0097       std::string parserCode;
0098       std::string parserPythonCode;  // for Python parsers
0099       if (type == kONNX) {
0100          // check first if we can load the SOFIE parser library
0101          if (gSystem->Load("libROOTTMVASofieParser") < 0) {
0102             throw std::runtime_error("RSofieReader: cannot use SOFIE with ONNX since libROOTTMVASofieParser is missing");
0103          }
0104          gInterpreter->Declare("#include \"TMVA/RModelParser_ONNX.hxx\"");
0105          parserCode += "{\nTMVA::Experimental::SOFIE::RModelParser_ONNX parser ; \n";
0106          if (verbose == 2)
0107             parserCode += "TMVA::Experimental::SOFIE::RModel model = parser.Parse(\"" + path + "\",true); \n";
0108          else
0109             parserCode += "TMVA::Experimental::SOFIE::RModel model = parser.Parse(\"" + path + "\"); \n";
0110       }
0111       else if (type == kKeras) {
0112          // use Keras Python parser
0113          parserPythonCode += "\"\"\"\n";
0114          parserPythonCode += "import ROOT\n";
0115
0116          // assume batch size is first entry in first input otherwise set to 1
0117          std::string batch_size = "1"; // need to fix parser with parm batch sizes
0118          if (!inputShapes.empty() && ! inputShapes[0].empty())
0119             batch_size = std::to_string(inputShapes[0][0]);
0120          parserPythonCode += "model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse('" + path + "'," + batch_size + ")\n";
0121       }
0122       else if (type == kPt) {
0123          // use PyTorch direct parser
0124          if (gSystem->Load("libROOTTMVASofiePyParsers") < 0) {
0125             throw std::runtime_error("RSofieReader: cannot use SOFIE with PyTorch since libROOTTMVASofiePyParsers is missing");
0126          }
0127          if (inputShapes.size() == 0) {
0128             throw std::runtime_error("RSofieReader: cannot use SOFIE with PyTorch since the input tensor shape is missing and is needed by the PyTorch parser");
0129          }
0130          std::string inputShapesStr = "{";
0131          for (unsigned int i = 0; i < inputShapes.size(); i++) {
0132             inputShapesStr += "{ ";
0133             for (unsigned int j = 0; j < inputShapes[i].size(); j++) {
0134                inputShapesStr += ROOT::Math::Util::ToString(inputShapes[i][j]);
0135                if (j < inputShapes[i].size()-1) inputShapesStr += ", ";
0136             }
0137             inputShapesStr += "}";
0138             if (i < inputShapes.size()-1) inputShapesStr += ", ";
0139          }
0140          inputShapesStr += "}";
0141          parserCode += "{\nTMVA::Experimental::SOFIE::RModel model = TMVA::Experimental::SOFIE::PyTorch::Parse(\"" + path + "\", "
0142                     + inputShapesStr + "); \n";
0143       }
0144       else if (type == kROOT) {
0145          // use  parser from ROOT
0146          parserCode += "{\nauto fileRead = TFile::Open(\"" + path + "\",\"READ\");\n";
0147          parserCode += "TMVA::Experimental::SOFIE::RModel * modelPtr;\n";
0148          parserCode += "auto keyList = fileRead->GetListOfKeys(); TString name;\n";
0149          parserCode += "for (const auto&& k : *keyList)  { \n";
0150          parserCode += "   TString cname =  ((TKey*)k)->GetClassName();  if (cname==\"TMVA::Experimental::SOFIE::RModel\") name = k->GetName(); }\n";
0151          parserCode += "fileRead->GetObject(name,modelPtr); fileRead->Close(); delete fileRead;\n";
0152          parserCode += "TMVA::Experimental::SOFIE::RModel & model = *modelPtr;\n";
0153       }
0154
0155        // add custom operators if needed
0156       if (fCustomOperators.size() > 0) {
0157          if (!parserPythonCode.empty())
0158             throw std::runtime_error("Cannot use Custom operator with a Python parser (e.g. from a Keras model)");
0159
0160          for (auto & op : fCustomOperators) {
0161             parserCode += "{ auto p = new TMVA::Experimental::SOFIE::ROperator_Custom<float>(\""
0162                       + op.fOpName + "\"," + op.fInputNames + "," + op.fOutputNames + "," + op.fOutputShapes + ",\"" + op.fFileName + "\");\n";
0163             parserCode += "std::unique_ptr<TMVA::Experimental::SOFIE::ROperator> op(p);\n";
0164             parserCode += "model.AddOperator(std::move(op));\n}\n";
0165          }
0166       }
0167
0168       int batchSize = 1;
0169       if (inputShapes.size() > 0 && inputShapes[0].size() > 0) {
0170          batchSize = inputShapes[0][0];
0171          if (batchSize < 1) batchSize = 1;
0172       }
0173       if (verbose) std::cout << "generating the code with batch size = " << batchSize << " ...\n";
0174
0175       if (parserPythonCode.empty()) {
0176          parserCode += "model.Generate(TMVA::Experimental::SOFIE::Options::kDefault,"
0177                     + ROOT::Math::Util::ToString(batchSize) + ", 0, " + std::to_string(verbose) + ");\n";
0178
0179          parserCode += "model.OutputGenerated(\"" + modelHeader + "\");\n";
0180          if (verbose) {
0181             parserCode += "model.PrintRequiredInputTensors();\n";
0182             parserCode += "model.PrintIntermediateTensors();\n";
0183             parserCode += "model.PrintOutputTensors();\n";
0184             if (verbose > 1)
0185                parserCode += "model.PrintGenerated(); \n";
0186          }
0187
0188          // need information on number of inputs (assume output is 1)
0189          parserCode += "int nInputs = model.GetInputTensorNames().size();\n";
0190
0191          //end of parsing C++ code
0192          parserCode += "return nInputs;\n}\n";
0193       } else {
0194          // Python case
0195          parserPythonCode += "model.Generate(ROOT.TMVA.Experimental.SOFIE.Options.kDefault,"
0196                    + ROOT::Math::Util::ToString(batchSize) + ", 0, " + std::to_string(verbose) + ")\n";
0197
0198          parserPythonCode += "model.OutputGenerated('" + modelHeader + "');\n";
0199          if (verbose) {
0200             parserPythonCode += "model.PrintRequiredInputTensors()\n";
0201             parserPythonCode += "model.PrintIntermediateTensors()\n";
0202             parserPythonCode += "model.PrintOutputTensors()\n";
0203             if (verbose > 1)
0204                parserPythonCode += "model.PrintGenerated()\n";
0205          }
0206          // end of Python parsing code
0207          parserPythonCode += "\"\"\"";
0208       }
0209       // executing parsing and generating code
0210       int iret = -1;
0211       if (parserPythonCode.empty()) {
0212          if (verbose) {
0213             std::cout << "...ParserCode being executed...:\n";
0214             std::cout << parserCode << std::endl;
0215          }
0216          iret = gROOT->ProcessLine(parserCode.c_str());
0217          fNInputs = iret;
0218       } else {
0219          if (verbose) {
0220             std::cout << "executing python3 -c ......" << std::endl;
0221             std::cout << parserPythonCode << std::endl;
0222          }
0223          iret = gSystem->Exec(TString("python3 -c ") + TString(parserPythonCode.c_str()));
0224          fNInputs = 1;
0225          // need number of inputs from input shapes
0226          if (!inputShapes.empty()) fNInputs = inputShapes.size();
0227       }
0228
0229       if (iret < 0) {
0230          std::string msg = "RSofieReader: error processing the parser code: \n" + parserCode;
0231          throw std::runtime_error(msg);
0232       } else if (verbose) {
0233          std::cout << "Model Header file is generated!" << std::endl;
0234       }
0235       if (fNInputs > 3) {
0236          throw std::runtime_error("RSofieReader does not yet support model with > 3 inputs");
0237       }
0238
0239       // compile now the generated code and create Session class
0240       if (verbose) std::cout << "compile generated code from file " <<modelHeader << std::endl;
0241       if (gSystem->AccessPathName(modelHeader.c_str())) {
0242          std::string msg = "RSofieReader: input header file " + modelHeader + " is not existing";
0243          throw std::runtime_error(msg);
0244       }
0245       if (verbose) std::cout << "Creating Inference function for model " << modelName << std::endl;
0246       std::string declCode;
0247       declCode += "#pragma cling optimize(2)\n";
0248       declCode += "#include \"" + modelHeader + "\"\n";
0249       // create global session instance: use UUID to have an unique name
0250       std::string sessionClassName = "TMVA_SOFIE_" + modelName + "::Session";
0251       TUUID uuid;
0252       std::string uidName = uuid.AsString();
0253       uidName.erase(std::remove_if(uidName.begin(), uidName.end(),
0254          []( char const& c ) -> bool { return !std::isalnum(c); } ), uidName.end());
0255
0256       std::string sessionName = "session_" + uidName;
0257       declCode += sessionClassName + " " + sessionName + "(\"" + modelWeights + "\");";
0258
0259       if (verbose) std::cout << "//global session declaration\n" << declCode << std::endl;
0260
0261       // need to load the ROOTTMVASOFIE library for some symbols used in generated code
0262       iret = gSystem->Load("libROOTTMVASofie");
0263       if (iret < 0)
0264          throw std::runtime_error("Error loading libROOTTMVASofie library");
0265
0266       bool ret = gInterpreter->Declare(declCode.c_str());
0267       if (!ret) {
0268          std::string msg = "RSofieReader: error compiling inference code and creating session class\n" + declCode;
0269          throw std::runtime_error(msg);
0270       }
0271
0272       fSessionPtr = (void *) gInterpreter->Calc(sessionName.c_str());
0273
0274       // define a function to be called for inference
0275       std::stringstream ifuncCode;
0276       std::string funcName = "SofieInference_" + uidName;
0277       ifuncCode << "std::vector<float> " + funcName + "( void * ptr";
0278       for (int i = 0; i < fNInputs; i++)
0279          ifuncCode << ", float * data" << i;
0280       ifuncCode << ") {\n";
0281       ifuncCode << "   " << sessionClassName << " * s = " << "(" << sessionClassName << "*) (ptr);\n";
0282       ifuncCode << "   return s->infer(";
0283       for (int i = 0; i < fNInputs; i++) {
0284          if (i>0) ifuncCode << ",";
0285          ifuncCode << "data" << i;
0286       }
0287       ifuncCode << ");\n";
0288       ifuncCode << "}\n";
0289
0290       if (verbose) std::cout << "//Inference function code using global session instance\n"
0291                               << ifuncCode.str() << std::endl;
0292
0293       ret = gInterpreter->Declare(ifuncCode.str().c_str());
0294       if (!ret) {
0295          std::string msg = "RSofieReader: error compiling inference function\n" + ifuncCode.str();
0296          throw std::runtime_error(msg);
0297       }
0298       fFuncPtr = (void *) gInterpreter->Calc(funcName.c_str());
0299       //fFuncPtr = reinterpret_cast<std::vector<float> (*)(void *, const float *)>(fptr);
0300       fInitialized = true;
0301    }
0302
0303    // Add custom operator
0304     void AddCustomOperator(const std::string &opName, const std::string &inputNames, const std::string & outputNames,
0305       const std::string & outputShapes, const std::string & fileName) {
0306          if (fInitialized)  std::cout << "WARNING: Model is already loaded and initialised. It must be done after adding the custom operators" << std::endl;
0307          fCustomOperators.push_back( {fileName, opName,inputNames, outputNames,outputShapes});
0308       }
0309
0310    // implementations for different outputs
0311    std::vector<float> DoCompute(const std::vector<float> & x1) {
0312       if (fNInputs != 1) {
0313          std::string msg = "Wrong number of inputs - model requires " + std::to_string(fNInputs);
0314          throw std::runtime_error(msg);
0315       }
0316       auto fptr = reinterpret_cast<std::vector<float> (*)(void *, const float *)>(fFuncPtr);
0317       return fptr(fSessionPtr, x1.data());
0318    }
0319    std::vector<float> DoCompute(const std::vector<float> & x1, const std::vector<float> & x2) {
0320       if (fNInputs != 2) {
0321          std::string msg = "Wrong number of inputs - model requires " + std::to_string(fNInputs);
0322          throw std::runtime_error(msg);
0323       }
0324       auto fptr = reinterpret_cast<std::vector<float> (*)(void *, const float *, const float *)>(fFuncPtr);
0325       return fptr(fSessionPtr, x1.data(),x2.data());
0326    }
0327    std::vector<float> DoCompute(const std::vector<float> & x1, const std::vector<float> & x2, const std::vector<float> & x3) {
0328       if (fNInputs != 3) {
0329          std::string msg = "Wrong number of inputs - model requires " + std::to_string(fNInputs);
0330          throw std::runtime_error(msg);
0331       }
0332       auto fptr = reinterpret_cast<std::vector<float> (*)(void *, const float *, const float *, const float *)>(fFuncPtr);
0333       return fptr(fSessionPtr, x1.data(),x2.data(),x3.data());
0334    }
0335
0336    /// Compute model prediction on vector
0337    template<typename... T>
0338    std::vector<float> Compute(T... x)
0339    {
0340       if(!fInitialized) {
0341          return std::vector<float>();
0342       }
0343
0344       // Take lock to protect model evaluation
0345       R__WRITE_LOCKGUARD(ROOT::gCoreMutex);
0346
0347       // Evaluate TMVA model (need to add support for multiple outputs)
0348       return DoCompute(x...);
0349
0350    }
0351    std::vector<float> Compute(const std::vector<float> &x) {
0352       if(!fInitialized) {
0353          return std::vector<float>();
0354       }
0355
0356       // Take lock to protect model evaluation
0357       R__WRITE_LOCKGUARD(ROOT::gCoreMutex);
0358
0359       // Evaluate TMVA model (need to add support for multiple outputs)
0360       return DoCompute(x);
0361    }
0362    /// Compute model prediction on input RTensor
0363    /// The shape of the input tensor should be {nevents, nfeatures}
0364    /// and the return shape will be {nevents, noutputs}
0365    /// support for now only a single input
0366    RTensor<float> Compute(RTensor<float> &x)
0367    {
0368       if(!fInitialized) {
0369          return RTensor<float>({0});
0370       }
0371       const auto nrows = x.GetShape()[0];
0372       const auto rowsize = x.GetStrides()[0];
0373       auto fptr = reinterpret_cast<std::vector<float> (*)(void *, const float *)>(fFuncPtr);
0374       auto result = fptr(fSessionPtr, x.GetData());
0375
0376       RTensor<float> y({nrows, result.size()}, MemoryLayout::ColumnMajor);
0377       std::copy(result.begin(),result.end(), y.GetData());
0378       //const bool layout = x.GetMemoryLayout() == MemoryLayout::ColumnMajor ? false : true;
0379       // assume column major layout
0380       for (size_t i = 1; i < nrows; i++) {
0381          result = fptr(fSessionPtr, x.GetData() + i*rowsize);
0382          std::copy(result.begin(),result.end(), y.GetData() + i*result.size());
0383       }
0384       return y;
0385    }
0386
0387 private:
0388
0389    bool fInitialized = false;
0390    int fNInputs = 0;
0391    void * fSessionPtr = nullptr;
0392    void * fFuncPtr = nullptr;
0393
0394    // data to insert custom operators
0395    struct CustomOperatorData {
0396       std::string fFileName; // code implementing the custom operator
0397       std::string fOpName; // operator name
0398       std::string fInputNames;  // input tensor names (convert as string as {"n1", "n2"})
0399       std::string fOutputNames;  // output tensor names converted as trind
0400       std::string fOutputShapes; // output shapes
0401    };
0402    std::vector<CustomOperatorData> fCustomOperators;
0403
0404 };
0405
0406 } // namespace Experimental
0407 } // namespace TMVA
0408
0409 #endif // TMVA_RREADER