Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
TMVA_SOFIE_PyTorch.C File Reference

Detailed Description

View in nbviewer Open in SWAN
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;
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";
//Running the Python script to generate PyTorch .pt file
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";
//Printing initialized tensors (weights)
std::cout<<"\n\n";
//Printing intermediate tensors
std::cout<<"\n\n";
//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\": ";
std::cout<<SOFIE::ConvertTypeToString(tensorType);
//Printing generated inference code
std::cout<<"\n\n";
model.PrintGenerated();
}
ROOT::Detail::TRangeCast< T, true > TRangeDynCast
TRangeDynCast is an adapter class that allows the typed iteration through a TCollection.
R__EXTERN TSystem * gSystem
Definition TSystem.h:572
bool CheckIfTensorAlreadyExist(std::string tensor_name)
Definition RModel.cxx:95
void OutputGenerated(std::string filename="", bool append=false)
Definition RModel.cxx:1284
const ETensorType & GetTensorType(std::string name) const
Definition RModel.cxx:67
const std::vector< size_t > & GetTensorShape(std::string name) const
Definition RModel.cxx:29
void Generate(std::underlying_type_t< Options > options, int batchSize=-1, long pos=0, bool verbose=false)
Definition RModel.cxx:916
static void PyInitialize()
Initialize Python interpreter.
Class supporting a collection of lines with C++ code.
Definition TMacro.h:31
Basic string class.
Definition TString.h:139
virtual Int_t Exec(const char *shellcmd)
Execute a command.
Definition TSystem.cxx:653
TString Python_Executable()
Function to find current Python executable used by ROOT If "Python3" is installed,...
TMarker m
Definition textangle.C:8
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.