root/TMVA/RModelParser_PyTorch.h

0001 // @(#)root/tmva/pymva $Id$
0002 // Author: Sanjiban Sengupta, 2021
0003
0004 /**********************************************************************************
0005  * Project: TMVA - a Root-integrated toolkit for multivariate data analysis       *
0006  * Package: TMVA                                                                  *
0007  *                                             *
0008  *                                                                                *
0009  * Description:                                                                   *
0010  *      Functionality for parsing a saved PyTorch .PT model into RModel object    *
0011  *                                                                                *
0012  * Authors (alphabetical):                                                        *
0013  *      Sanjiban Sengupta <sanjiban.sg@gmail.com>                                 *
0014  *                                                                                *
0015  * Copyright (c) 2021:                                                            *
0016  *      CERN, Switzerland                                                         *
0017  *                                                                                *
0018  *                                                                                *
0019  * Redistribution and use in source and binary forms, with or without             *
0020  * modification, are permitted according to the terms listed in LICENSE           *
0021  * (see tmva/doc/LICENSE)                                          *
0022  **********************************************************************************/
0023
0024
0025 #ifndef TMVA_SOFIE_RMODELPARSER_PYTORCH
0026 #define TMVA_SOFIE_RMODELPARSER_PYTORCH
0027
0028 #include "TMVA/RModel.hxx"
0029 #include "TMVA/SOFIE_common.hxx"
0030 #include "TMVA/Types.h"
0031 #include "TMVA/OperatorList.hxx"
0032
0033 #include "Rtypes.h"
0034 #include "TString.h"
0035
0036
0037 namespace TMVA::Experimental::SOFIE::PyTorch {
0038
0039 /// Parser function for translating PyTorch .pt model into a RModel object.
0040 /// Accepts the file location of a PyTorch model, shapes and data-types of input tensors
0041 /// and returns the equivalent RModel object.
0042 RModel Parse(std::string filepath, std::vector<std::vector<size_t>> inputShapes, std::vector<ETensorType> dtype);
0043
0044 /// Overloaded Parser function for translating PyTorch .pt model into a RModel object.
0045 /// Accepts the file location of a PyTorch model and the shapes of input tensors.
0046 /// Builds the vector of data-types for input tensors and calls the `Parse()` function to
0047 /// return the equivalent RModel object.
0048 RModel Parse(std::string filepath, std::vector<std::vector<size_t>> inputShapes);
0049
0050 } // namespace TMVA::Experimental::SOFIE::PyTorch
0051
0052 #endif //TMVA_PYMVA_RMODELPARSER_PYTORCH