RModelParser_ONNX.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_RMODELPARSER_ONNX
#define TMVA_SOFIE_RMODELPARSER_ONNX
 
#include "TMVA/RModel.hxx"
 
#include <memory>
#include <functional>
#include <unordered_map>
#include <fstream>
 
// forward declaration
namespace onnx {
class NodeProto;
class GraphProto;
class ModelProto;
class TensorProto;
} // namespace onnx
 
namespace TMVA {
namespace Experimental {
namespace SOFIE {
 
class RModelParser_ONNX;
 
using ParserFuncSignature =
   std::function<std::unique_ptr<ROperator>(RModelParser_ONNX & /*parser*/, const onnx::NodeProto & /*nodeproto*/)>;
using ParserFuseFuncSignature =
   std::function<std::unique_ptr<ROperator> (RModelParser_ONNX& /*parser*/, const onnx::NodeProto& /*firstnode*/, const onnx::NodeProto& /*secondnode*/)>;
 
class RModelParser_ONNX {
public:
   struct OperatorsMapImpl;
 
   enum EFusedOp { kMatMulAdd, kConvAdd, kConvTransAdd, kGemmRelu, kBatchnormRelu};
 
private:
 
   bool fVerbose = false;
   // Registered operators
   std::unique_ptr<OperatorsMapImpl> fOperatorsMapImpl;
   // Type of the tensors
   std::unordered_map<std::string, ETensorType> fTensorTypeMap;
 
   // List of fused operators storing as key the second operator and a value a pair of fusion type and parent operator
   std::map<int, std::pair<EFusedOp, int>> fFusedOperators;
 
   //  weight data file
   std::ifstream fDataFile;
   // filename of model
   std::string fDataFileName;
 
 
public:
   // Register an ONNX operator
   void RegisterOperator(const std::string &name, ParserFuncSignature func);
 
   // Check if the operator is registered
   bool IsRegisteredOperator(const std::string &name);
 
   // List of registered operators (in alphabetical order)
   std::vector<std::string> GetRegisteredOperators();
 
   // Set the type of the tensor
   void RegisterTensorType(const std::string & /*name*/, ETensorType /*type*/);
 
   // Check if the type of the tensor is registered
   bool IsRegisteredTensorType(const std::string & /*name*/);
 
   // check verbosity
   bool Verbose() const {
      return fVerbose;
   }
 
   // Get the type of the tensor
   ETensorType GetTensorType(const std::string &name);
 
   // Parse the index'th node from the ONNX graph
   std::unique_ptr<ROperator> ParseOperator(const size_t /*index*/, const onnx::GraphProto & /*graphproto*/,
                                            const std::vector<size_t> & /*nodes*/, const std::vector<int> & /* children */);
 
   // check a graph for missing operators
   void CheckGraph(const onnx::GraphProto & g, int & level, std::map<std::string, int> & missingOperators);
 
   // parse the ONNX graph
   void ParseONNXGraph(RModel & model, const onnx::GraphProto & g, std::string  name = "");
 
   std::unique_ptr<onnx::ModelProto> LoadModel(const std::string &filename);
   std::unique_ptr<onnx::ModelProto> LoadModel(std::istream &input);
 
   std::shared_ptr<void> GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_length, ETensorType type );
 
public:
 
   RModelParser_ONNX() noexcept;
 
   RModel Parse(std::string const &filename, bool verbose = false);
   RModel Parse(std::istream &input, std::string const &name, bool verbose = false);
 
   // check the model for missing operators - return false in case some operator implementation is missing
   bool CheckModel(std::string filename, bool verbose = false);
 
   //set external data full path (needed if external data are not stored in the default modelName.onnx.data)
   // call this function before parsing
   void SetExternalDataFile(const std::string & dataFileName) {
      fDataFileName = dataFileName;
   }
 
   ~RModelParser_ONNX();
};
 
} // namespace SOFIE
} // namespace Experimental
} // namespace TMVA
 
#endif // TMVA_SOFIE_RMODELPARSER_ONNX