TMVA_SOFIE_PyTorch.py 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.

 
import sys
 
import ROOT
import torch
import torch.nn as nn
 
SOFIE = ROOT.TMVA.Experimental.SOFIE
 
# Python and C++ write to separate stdout buffers; flush both on every line so
# that the Python prints and the RModel printouts appear in order
sys.stdout.reconfigure(line_buffering=True)
ROOT.gInterpreter.ProcessLine("std::cout << std::unitbuf;")
 
# ------------------------------------------------------------------------------
# Step 1: Create, train and save a simple PyTorch model
# ------------------------------------------------------------------------------
 
model = nn.Sequential(
    nn.Linear(32, 16),
    nn.ReLU(),
    nn.Linear(16, 8),
    nn.ReLU(),
)
 
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
 
x = torch.randn(2, 32)
y = torch.randn(2, 8)
 
for i in range(500):
    y_pred = model(x)
    loss = criterion(y_pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
 
model.eval()
m = torch.jit.script(model)
torch.jit.save(m, "PyTorchModel.pt")
 
# ------------------------------------------------------------------------------
# Step 2: Parse the saved PyTorch .pt file with TMVA::SOFIE
# ------------------------------------------------------------------------------
 
# Parsing a PyTorch model requires the shape and data-type of input tensor
# Data-type of input tensor defaults to Float if not specified
input_shapes = ROOT.std.vector["std::vector<std::size_t>"]()
input_shapes.push_back([2, 32])
 
# Parsing the saved PyTorch .pt file into RModel object
model = SOFIE.PyTorch.Parse("PyTorchModel.pt", input_shapes)
 
# Generating inference code
model.Generate()
model.OutputGenerated("PyTorchModel.hxx")
 
# Printing required input tensors
print("\n")
model.PrintRequiredInputTensors()
 
# Printing initialized tensors (weights)
print("\n")
model.PrintInitializedTensors()
 
# Printing intermediate tensors
print("\n")
model.PrintIntermediateTensors()
 
# Checking if tensor already exist in model
tensor_exists = bool(model.CheckIfTensorAlreadyExist("0weight"))
print(f'\n\nTensor "0weight" already exist: {str(tensor_exists).lower()}\n')
 
tensor_shape = model.GetTensorShape("0weight")
print('Shape of tensor "0weight": ' + ",".join(str(dim) for dim in tensor_shape) + ",")
 
tensor_type = model.GetTensorType("0weight")
print(f'\nData type of tensor "0weight": {SOFIE.ConvertTypeToString(tensor_type)}')
 
# Printing generated inference code
print()
model.PrintGenerated()

Torch Version: 2.13.0+cpu
 
 
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: "result"   type: float shape: [2,16]
Tensor name: "input2"   type: float shape: [2,8]
Tensor name: "input0"   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 [Tue Jul 14 12:50:01 202] 
 
#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;
inline void doInfer(Session const &session, float const* tensor_input1, float *tensor_result3 );
struct Session {
// initialized (weights and constant) 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>(256);
 
 
// --- 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() + 64);
 
 // Allocating memory for intermediate tensor result3 with size 64 bytes
float* tensor_result3 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 0);
 
 
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();
 
}
 
 
 
std::vector<float> infer(float const* tensor_input1){
   std::vector<float > output_tensor_result3(16);
   doInfer(*this, tensor_input1, output_tensor_result3.data() );
   return {output_tensor_result3};
}
};   // end of Session
 
 
// Input tensor dimensions
using TMVA::Experimental::SOFIE::SingleDim;
using TMVA::Experimental::SOFIE::TensorDims;
using TMVA::Experimental::SOFIE::makeDims;
 
constexpr std::array<SingleDim, 2> dim_input1{SingleDim{2}, SingleDim{32}};
 
constexpr std::array<TensorDims, 1> inputTensorDims{
   makeDims(dim_input1)
};
 
constexpr bool hasDynamicInputTensors{false};
 
 
// Output tensor dimensions
constexpr std::array<SingleDim, 2> dim_result3{SingleDim{2}, SingleDim{8}};
 
constexpr std::array<TensorDims, 1> outputTensorDims{
   makeDims(dim_result3)
};
 
constexpr bool hasDynamicOutputTensors{false};
 
inline void doInfer(Session const &session, float const* tensor_input1, float *tensor_result3 ) {
 
    auto &tensor_0bias = session.tensor_0bias;
    auto &tensor_0weight = session.tensor_0weight;
    auto &tensor_2bias = session.tensor_2bias;
    auto &tensor_2weight = session.tensor_2weight;
    auto &tensor_input0 = session.tensor_input0;
    auto &tensor_input2 = session.tensor_input2;
    auto &tensor_result = session.tensor_result;
 
 
//--------- Gemm op_0 { 2 , 32 } * { 16 , 32 } -> { 2 , 16 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 16 * j;
      TMVA::Experimental::SOFIE::Copy(tensor_input0 + y_index, tensor_0bias, 16);
   }
   TMVA::Experimental::SOFIE::Gemm_Call(tensor_input0, true, false, 16, 2, 32, 1, tensor_0weight, tensor_input1, 1,nullptr);
 
//------ RELU
   for (int id = 0; id < 32 ; id++){
      tensor_result[id] = ((tensor_input0[id] > 0 )? tensor_input0[id] : 0);
   }
 
//--------- Gemm op_2 { 2 , 16 } * { 8 , 16 } -> { 2 , 8 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 8 * j;
      TMVA::Experimental::SOFIE::Copy(tensor_input2 + y_index, tensor_2bias, 8);
   }
   TMVA::Experimental::SOFIE::Gemm_Call(tensor_input2, true, false, 8, 2, 16, 1, tensor_2weight, tensor_result, 1,nullptr);
 
//------ RELU
   for (int id = 0; id < 16 ; id++){
      tensor_result3[id] = ((tensor_input2[id] > 0 )? tensor_input2[id] : 0);
   }
 
}
} //TMVA_SOFIE_PyTorchModel
 
#endif  // ROOT_TMVA_SOFIE_PYTORCHMODEL

Author

Sanjiban Sengupta

Definition in file TMVA_SOFIE_PyTorch.py.