Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
TMVA_SOFIE_RDataFrame.py File Reference

Detailed Description

View in nbviewer Open in SWAN
Example of inference with SOFIE and RDataFrame, of a model trained with PyTorch.

First, generate the input ONNX model by running TMVA_SOFIE_PyTorch_HiggsModel.py.

This tutorial parses the input model and runs the inference using ROOT's JITing capability.

from os.path import exists
import ROOT
# check if the input file exists
modelFile = "HiggsModel.onnx"
modelName = "HiggsModel"
if not exists(modelFile):
raise FileNotFoundError("You need to run TMVA_SOFIE_PyTorch_HiggsModel.py to generate the ONNX trained model")
# parse the input ONNX model into RModel object
model = parser.Parse(modelFile)
# generating inference code
model.OutputGenerated("Higgs_trained_model_generated.hxx")
# compile using ROOT JIT trained model
print("compiling SOFIE model and functor....")
ROOT.gInterpreter.Declare('#include "Higgs_trained_model_generated.hxx"')
ROOT.gInterpreter.Declare('auto sofie_functor = TMVA::Experimental::SofieFunctor<7,TMVA_SOFIE_'+modelName+'::Session>(0,"Higgs_trained_model_generated.dat");')
# run inference over input data
inputFile = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
df1 = ROOT.RDataFrame("sig_tree", inputFile)
h1 = df1.Define("DNN_Value", "sofie_functor(rdfslot_,m_jj, m_jjj, m_lv, m_jlv, m_bb, m_wbb, m_wwbb)").Histo1D(("h_sig", "", 100, 0, 1),"DNN_Value")
df2 = ROOT.RDataFrame("bkg_tree", inputFile)
h2 = df2.Define("DNN_Value", "sofie_functor(rdfslot_,m_jj, m_jjj, m_lv, m_jlv, m_bb, m_wbb, m_wwbb)").Histo1D(("h_bkg", "", 100, 0, 1),"DNN_Value")
# run over the input data once, combining both RDataFrame graphs.
print("Number of signal entries",h1.GetEntries())
print("Number of background entries",h2.GetEntries())
h1.DrawClone("SAME")
ROOT::Detail::TRangeCast< T, true > TRangeDynCast
TRangeDynCast is an adapter class that allows the typed iteration through a TCollection.
ROOT's RDataFrame offers a modern, high-level interface for analysis of data stored in TTree ,...
//Code generated automatically by TMVA for Inference of Model file [HiggsModel.onnx] at [Sun Aug 2 02:36:17 2026]
#ifndef ROOT_TMVA_SOFIE_HIGGSMODEL
#define ROOT_TMVA_SOFIE_HIGGSMODEL
#include <cassert>
#include <cmath>
#include <vector>
#include <cstdint>
#include <limits>
#include <algorithm>
#include <iomanip>
#include <cstring>
#include <map>
#include <sstream>
#include <string>
#include <memory>
#include <iostream>
#include <stdexcept>
#include <algorithm>
#include <array>
#include <cstddef>
#include <istream>
#include <limits>
#include <stdexcept>
#include <string>
#include <string_view>
#include <fstream>
namespace TMVA_SOFIE_HiggsModel{
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
// --- Standalone SOFIE inference helper functions ---
inline void Gemm_Call(float *output, bool transa, bool transb, int m, int n, int k, float alpha, const float *A,
const float *B, float beta, const float *C)
{
char ct = 't';
char cn = 'n';
const int *lda = transa ? &k : &m;
const int *ldb = transb ? &n : &k;
const int *ldc = &m;
if (C != nullptr) {
std::copy(C, C + m * n, output);
}
BLAS::sgemm_(transa ? &ct : &cn, transb ? &ct : &cn, &m, &n, &k, &alpha, A, lda, B, ldb, &beta, output, ldc);
}
inline void Relu(float *output, float const *input, int size)
{
for (int i = 0; i < size; i++) {
output[i] = (input[i] > 0.0f) ? input[i] : 0.0f;
}
}
inline void Fill(float *output, float value, int size)
{
std::fill(output, output + size, value);
}
template <class T>
inline void Copy(T *output, T const *input, int size)
{
std::copy(input, input + size, output);
}
inline float ParseFloatToken(const std::string &s)
{
if (s == "inf")
return std::numeric_limits<float>::infinity();
if (s == "-inf")
return -std::numeric_limits<float>::infinity();
if (s == "nan")
return std::numeric_limits<float>::quiet_NaN();
return std::stof(s);
}
template <class T>
void ReadTensorFromStream(std::istream &is, T &target, std::string const &expectedName, std::size_t expectedLength)
{
std::string name;
std::size_t length;
is >> name >> length;
if (name != expectedName) {
std::string err_msg =
"TMVA-SOFIE failed to read the correct tensor name; expected name is " + expectedName + " , read " + name;
throw std::runtime_error(err_msg);
}
if (length != expectedLength) {
std::string err_msg = "TMVA-SOFIE failed to read the correct tensor size; expected size is " +
std::to_string(expectedLength) + " , read " + std::to_string(length);
throw std::runtime_error(err_msg);
}
std::string token;
for (std::size_t i = 0; i < length; ++i) {
is >> token;
target[i] = ParseFloatToken(token);
}
if (is.fail()) {
throw std::runtime_error("TMVA-SOFIE failed to read the values for tensor " + expectedName);
}
}
struct SingleDim {
enum class Kind { Static, Symbolic };
Kind kind;
std::size_t dim;
std::string_view name;
constexpr SingleDim(std::size_t v) : kind(Kind::Static), dim(v), name() {}
constexpr SingleDim(const char *v) : kind(Kind::Symbolic), dim(0), name(v) {}
};
struct TensorDims {
const SingleDim *data;
std::size_t size;
constexpr std::size_t total_size() const
{
std::size_t result = 1;
for (std::size_t i = 0; i < size; ++i) {
result *= data[i].dim;
}
return result;
}
};
template <class Arr>
constexpr TensorDims makeDims(Arr const &arr)
{
return TensorDims{arr.data(), arr.size()};
}
// --- End of SOFIE inference helper functions ---
struct Session;
inline void doInfer(Session const &session, float const* tensor_input, float *tensor_output );
struct Session {
// initialized (weights and constant) tensors
std::vector<float> fTensor_4weight = std::vector<float>(64);
float * tensor_4weight = fTensor_4weight.data();
std::vector<float> fTensor_2bias = std::vector<float>(64);
float * tensor_2bias = fTensor_2bias.data();
std::vector<float> fTensor_4bias = std::vector<float>(1);
float * tensor_4bias = fTensor_4bias.data();
std::vector<float> fTensor_2weight = std::vector<float>(4096);
float * tensor_2weight = fTensor_2weight.data();
std::vector<float> fTensor_0bias = std::vector<float>(64);
float * tensor_0bias = fTensor_0bias.data();
std::vector<float> fTensor_0weight = std::vector<float>(448);
float * tensor_0weight = fTensor_0weight.data();
//--- Allocating session memory pool to be used for allocating intermediate tensors
std::vector<char> fIntermediateMemoryPool = std::vector<char>(512);
// --- Positioning intermediate tensor memory --
// Allocating memory for intermediate tensor relu with size 256 bytes
float* tensor_relu = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 0);
// Allocating memory for intermediate tensor relu_1 with size 256 bytes
float* tensor_relu_1 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 256);
// Allocating memory for intermediate tensor linear_2 with size 4 bytes
float* tensor_linear_2 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 252);
// Allocating memory for intermediate tensor output with size 4 bytes
float* tensor_output = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 248);
Session(std::string filename ="HiggsModel.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");
}
ReadTensorFromStream(f, tensor_4weight, "tensor_4weight", 64);
ReadTensorFromStream(f, tensor_2bias, "tensor_2bias", 64);
ReadTensorFromStream(f, tensor_4bias, "tensor_4bias", 1);
ReadTensorFromStream(f, tensor_2weight, "tensor_2weight", 4096);
ReadTensorFromStream(f, tensor_0bias, "tensor_0bias", 64);
ReadTensorFromStream(f, tensor_0weight, "tensor_0weight", 448);
f.close();
}
std::vector<float> infer(float const* tensor_input){
std::vector<float > output_tensor_output(1);
doInfer(*this, tensor_input, output_tensor_output.data() );
return {output_tensor_output};
}
}; // end of Session
// Input tensor dimensions
constexpr std::array<SingleDim, 2> dim_input{SingleDim{1}, SingleDim{7}};
constexpr std::array<TensorDims, 1> inputTensorDims{
makeDims(dim_input)
};
constexpr bool hasDynamicInputTensors{false};
// Output tensor dimensions
constexpr std::array<SingleDim, 2> dim_output{SingleDim{1}, SingleDim{1}};
constexpr std::array<TensorDims, 1> outputTensorDims{
makeDims(dim_output)
};
constexpr bool hasDynamicOutputTensors{false};
inline void doInfer(Session const &session, float const* tensor_input, float *tensor_output ) {
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_4bias = session.tensor_4bias;
auto &tensor_4weight = session.tensor_4weight;
auto &tensor_linear_2 = session.tensor_linear_2;
auto &tensor_relu = session.tensor_relu;
auto &tensor_relu_1 = session.tensor_relu_1;
//--------- Gemm op_0 { 1 , 7 } * { 64 , 7 } -> { 1 , 64 }
for (size_t j = 0; j < 1; j++) {
size_t y_index = 64 * j;
Copy(tensor_relu + y_index, tensor_0bias, 64);
}
Gemm_Call(tensor_relu, true, false, 64, 1, 7, 1, tensor_0weight, tensor_input, 1,nullptr);
//--- applying RELU to output
Relu(tensor_relu, tensor_relu, 64);
//--------- Gemm op_1 { 1 , 64 } * { 64 , 64 } -> { 1 , 64 }
for (size_t j = 0; j < 1; j++) {
size_t y_index = 64 * j;
Copy(tensor_relu_1 + y_index, tensor_2bias, 64);
}
Gemm_Call(tensor_relu_1, true, false, 64, 1, 64, 1, tensor_2weight, tensor_relu, 1,nullptr);
//--- applying RELU to output
Relu(tensor_relu_1, tensor_relu_1, 64);
//--------- Gemm op_2 { 1 , 64 } * { 1 , 64 } -> { 1 , 1 }
for (size_t j = 0; j < 1; j++) {
size_t y_index = j;
Copy(tensor_linear_2 + y_index, tensor_4bias, 1);
}
Gemm_Call(tensor_linear_2, true, false, 1, 1, 64, 1, tensor_4weight, tensor_relu_1, 1,nullptr);
//------ Sigmoid -- 3
for (int id = 0; id < 1 ; id++){
tensor_output[id] = 1 / (1 + std::exp( - tensor_linear_2[id]));
}
}
} //TMVA_SOFIE_HiggsModel
namespace clad {
namespace custom_derivatives {
namespace TMVA_SOFIE_HiggsModel {
using ::TMVA_SOFIE_HiggsModel::Gemm_Call;
inline void Gemm_Call_pullback(float *output, bool transa, bool transb, int m, int n, int k, float alpha,
const float *A, const float *B, float beta, const float *C, float *_d_output, bool *,
bool *, int *, int *, int *, float *_d_alpha, float *_d_A, float *_d_B, float *_d_beta,
float *_d_C)
{
// TODO:
// - fix and test the implementation for alpha != 1.0
if (alpha != 1.0f) {
return;
}
// beta needs to be one because we want to add to _d_A and _d_B instead of
// overwriting it.
float one = 1.;
// ---- dA ----
if (!transa) {
// dA += dY * op(B)^T
Gemm_Call(_d_A, false, !transb, m, k, n, one, _d_output, B, one, _d_A);
} else {
// dA += op(B) * dY^T
Gemm_Call(_d_A, transb, true, k, m, n, one, B, _d_output, one, _d_A);
}
// ---- dB ----
if (!transb) {
// dB += op(A)^T * dY
Gemm_Call(_d_B, !transa, false, k, n, m, one, A, _d_output, one, _d_B);
} else {
// dB += dY^T * op(A)
Gemm_Call(_d_B, true, transa, n, k, m, one, _d_output, A, one, _d_B);
}
int sizeC = n * m;
for (int i = 0; i < sizeC; ++i) {
if (C) {
*_d_alpha += _d_output[i] * (output[i] - beta * C[i]);
*_d_beta += _d_output[i] * C[i];
} else {
*_d_alpha += _d_output[i] * output[i];
}
if (_d_C)
_d_C[i] += _d_output[i] * beta;
}
}
inline void Copy_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *)
{
for (int i = 0; i < size; i++) {
output[i] = input[i];
_d_input[i] += _d_output[i];
_d_output[i] = 0.F;
}
}
inline void Fill_pullback(float *output, float value, int size, float *_d_output, float *_d_value, int *)
{
for (int i = 0; i < size; i++) {
output[i] = value;
*_d_value += _d_output[i];
_d_output[i] = 0.F;
}
}
inline void Relu_pullback(float *output, const float *input, int size, float *_d_output, float *_d_input, int *)
{
for (int i = 0; i < size; i++) {
output[i] = input[i] > 0.F ? input[i] : 0.F;
float _r_d0 = _d_output[i];
_d_output[i] = 0.F;
if (input[i] > 0.F)
_d_input[i] += _r_d0;
}
}
} // namespace TMVA_SOFIE_HiggsModel
} // namespace custom_derivatives
} // namespace clad
#endif // ROOT_TMVA_SOFIE_HIGGSMODEL
compiling SOFIE model and functor....
Number of signal entries 10000.0
Number of background entries 10000.0
Author
Lorenzo Moneta

Definition in file TMVA_SOFIE_RDataFrame.py.