TMVA_SOFIE_PyTorch.C File Reference

Detailed Description

This macro provides a simple example for the parsing of PyTorch .pt file into RModel object and further generating the .hxx header files for inference.

 
using namespace TMVA::Experimental;
 
TString pythonSrc = "\
import torch\n\
import torch.nn as nn\n\
\n\
model = nn.Sequential(\n\
           nn.Linear(32,16),\n\
           nn.ReLU(),\n\
           nn.Linear(16,8),\n\
           nn.ReLU()\n\
           )\n\
\n\
criterion = nn.MSELoss()\n\
optimizer = torch.optim.SGD(model.parameters(),lr=0.01)\n\
\n\
x=torch.randn(2,32)\n\
y=torch.randn(2,8)\n\
\n\
for i in range(500):\n\
    y_pred = model(x)\n\
    loss = criterion(y_pred,y)\n\
    optimizer.zero_grad()\n\
    loss.backward()\n\
    optimizer.step()\n\
\n\
model.eval()\n\
m = torch.jit.script(model)\n\
torch.jit.save(m,'PyTorchModel.pt')\n";
 
 
void TMVA_SOFIE_PyTorch(){
 
    //Running the Python script to generate PyTorch .pt file
    TMVA::PyMethodBase::PyInitialize();
 
    TMacro m;
    m.AddLine(pythonSrc);
    m.SaveSource("make_pytorch_model.py");
    gSystem->Exec(TMVA::Python_Executable() + " make_pytorch_model.py");
 
    //Parsing a PyTorch model requires the shape and data-type of input tensor
    //Data-type of input tensor defaults to Float if not specified
    std::vector<size_t> inputTensorShapeSequential{2,32};
    std::vector<std::vector<size_t>> inputShapesSequential{inputTensorShapeSequential};
 
    //Parsing the saved PyTorch .pt file into RModel object
    SOFIE::RModel model = SOFIE::PyTorch::Parse("PyTorchModel.pt",inputShapesSequential);
 
    //Generating inference code
    model.Generate();
    model.OutputGenerated("PyTorchModel.hxx");
 
    //Printing required input tensors
    std::cout<<"\n\n";
    model.PrintRequiredInputTensors();
 
    //Printing initialized tensors (weights)
    std::cout<<"\n\n";
    model.PrintInitializedTensors();
 
    //Printing intermediate tensors
    std::cout<<"\n\n";
    model.PrintIntermediateTensors();
 
    //Checking if tensor already exist in model
    std::cout<<"\n\nTensor \"0weight\" already exist: "<<std::boolalpha<<model.CheckIfTensorAlreadyExist("0weight")<<"\n\n";
    std::vector<size_t> tensorShape = model.GetTensorShape("0weight");
    std::cout<<"Shape of tensor \"0weight\": ";
    for(auto& it:tensorShape){
        std::cout<<it<<",";
    }
    std::cout<<"\n\nData type of tensor \"0weight\": ";
    SOFIE::ETensorType tensorType = model.GetTensorType("0weight");
    std::cout<<SOFIE::ConvertTypeToString(tensorType);
 
    //Printing generated inference code
    std::cout<<"\n\n";
    model.PrintGenerated();
}

 
Torch Version: 2.7.0+cu126
RecursiveScriptModule(
  original_name=Sequential
  (0): RecursiveScriptModule(original_name=Linear)
  (1): RecursiveScriptModule(original_name=ReLU)
  (2): RecursiveScriptModule(original_name=Linear)
  (3): RecursiveScriptModule(original_name=ReLU)
)
RecursiveScriptModule(
  original_name=Sequential
  (0): RecursiveScriptModule(original_name=Linear)
  (1): RecursiveScriptModule(original_name=ReLU)
  (2): RecursiveScriptModule(original_name=Linear)
  (3): RecursiveScriptModule(original_name=ReLU)
)
 
 
Model requires following inputs:
Fully Specified Tensor name: input1 type: float shape: [2,32]
 
 
 
Model initialized the following tensors:
Tensor name: "2bias" type: float shape: [8]
Tensor name: "0weight"  type: float shape: [16,32]
Tensor name: "2weight"  type: float shape: [8,16]
Tensor name: "0bias" type: float shape: [16]
 
 
 
Model specify the following intermediate tensors:
Tensor name: "result3"  type: float shape: [2,8]
Tensor name: "2biasbcast"  type: float shape: [2,8]
Tensor name: "input2"   type: float shape: [2,8]
Tensor name: "input0"   type: float shape: [2,16]
Tensor name: "result"   type: float shape: [2,16]
Tensor name: "0biasbcast"  type: float shape: [2,16]
 
 
 
Tensor "0weight" already exist: true
 
Shape of tensor "0weight": 16,32,
 
Data type of tensor "0weight": float
 
//Code generated automatically by TMVA for Inference of Model file [PyTorchModel.pt] at [Sat May 24 02:58:28 2025] 
 
#ifndef ROOT_TMVA_SOFIE_PYTORCHMODEL
#define ROOT_TMVA_SOFIE_PYTORCHMODEL
 
#include <algorithm>
#include <vector>
#include "TMVA/SOFIE_common.hxx"
#include <fstream>
 
namespace TMVA_SOFIE_PyTorchModel{
namespace BLAS{
   extern "C" void sgemv_(const char * trans, const int * m, const int * n, const float * alpha, const float * A,
                          const int * lda, const float * X, const int * incx, const float * beta, const float * Y, const int * incy);
   extern "C" void sgemm_(const char * transa, const char * transb, const int * m, const int * n, const int * k,
                          const float * alpha, const float * A, const int * lda, const float * B, const int * ldb,
                          const float * beta, float * C, const int * ldc);
}//BLAS
struct Session {
// initialized tensors
std::vector<float> fTensor_2bias = std::vector<float>(8);
float * tensor_2bias = fTensor_2bias.data();
std::vector<float> fTensor_0weight = std::vector<float>(512);
float * tensor_0weight = fTensor_0weight.data();
std::vector<float> fTensor_2weight = std::vector<float>(128);
float * tensor_2weight = fTensor_2weight.data();
std::vector<float> fTensor_0bias = std::vector<float>(16);
float * tensor_0bias = fTensor_0bias.data();
 
//--- Allocating session memory pool to be used for allocating intermediate tensors
std::vector<char> fIntermediateMemoryPool = std::vector<char>(384);
 
 
// --- Positioning intermediate tensor memory --
 // Allocating memory for intermediate tensor input0 with size 128 bytes
float* tensor_input0 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 0);
 
 // Allocating memory for intermediate tensor result with size 128 bytes
float* tensor_result = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 128);
 
 // Allocating memory for intermediate tensor input2 with size 64 bytes
float* tensor_input2 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 256);
 
 // Allocating memory for intermediate tensor result3 with size 64 bytes
float* tensor_result3 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 320);
 
//--- declare and allocate the intermediate tensors
std::vector<float> fTensor_2biasbcast = std::vector<float>(16);
float * tensor_2biasbcast = fTensor_2biasbcast.data();
std::vector<float> fTensor_0biasbcast = std::vector<float>(32);
float * tensor_0biasbcast = fTensor_0biasbcast.data();
 
 
Session(std::string filename ="PyTorchModel.dat") {
 
//--- reading weights from file
   std::ifstream f;
   f.open(filename);
   if (!f.is_open()) {
      throw std::runtime_error("tmva-sofie failed to open file " + filename + " for input weights");
   }
   using TMVA::Experimental::SOFIE::ReadTensorFromStream;
   ReadTensorFromStream(f, tensor_2bias, "tensor_2bias", 8);
   ReadTensorFromStream(f, tensor_0weight, "tensor_0weight", 512);
   ReadTensorFromStream(f, tensor_2weight, "tensor_2weight", 128);
   ReadTensorFromStream(f, tensor_0bias, "tensor_0bias", 16);
   f.close();
 
//--- broadcast bias tensor 0biasfor Gemm op
   {
      float * data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<float>(tensor_0bias,{ 16 }, { 2 , 16 });
      std::copy(data, data + 32, tensor_0biasbcast);
      delete [] data;
   }
//--- broadcast bias tensor 2biasfor Gemm op
   {
      float * data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<float>(tensor_2bias,{ 8 }, { 2 , 8 });
      std::copy(data, data + 16, tensor_2biasbcast);
      delete [] data;
   }
}
 
void doInfer(float const* tensor_input1,  std::vector<float> &output_tensor_result3 ){
 
 
//--------- Gemm
   TMVA::Experimental::SOFIE::Gemm_Call(tensor_input0, true, false, 16, 2, 32, 1,tensor_0weight, tensor_input1, 1,tensor_0biasbcast);
 
//------ RELU
   for (int id = 0; id < 32 ; id++){
      tensor_result[id] = ((tensor_input0[id] > 0 )? tensor_input0[id] : 0);
   }
 
//--------- Gemm
   TMVA::Experimental::SOFIE::Gemm_Call(tensor_input2, true, false, 8, 2, 16, 1,tensor_2weight, tensor_result, 1,tensor_2biasbcast);
 
//------ RELU
   for (int id = 0; id < 16 ; id++){
      tensor_result3[id] = ((tensor_input2[id] > 0 )? tensor_input2[id] : 0);
   }
   using TMVA::Experimental::SOFIE::UTILITY::FillOutput;
 
   FillOutput(tensor_result3, output_tensor_result3, 16);
}
 
 
 
std::vector<float> infer(float const* tensor_input1){
   std::vector<float > output_tensor_result3;
   doInfer(tensor_input1, output_tensor_result3 );
   return {output_tensor_result3};
}
};   // end of Session
 
} //TMVA_SOFIE_PyTorchModel
 
#endif  // ROOT_TMVA_SOFIE_PYTORCHMODEL

Author

Sanjiban Sengupta

Definition in file TMVA_SOFIE_PyTorch.C.