TMVA_SOFIE_Keras.py File Reference

Detailed Description

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

 
 
import contextlib
import warnings
 
import numpy as np
import ROOT
from tensorflow.keras.layers import Activation, Dense, Input, Softmax
from tensorflow.keras.models import Model
 
# Enable ROOT in batch mode (same effect as -nodraw)
ROOT.gROOT.SetBatch(True)
 
 
@contextlib.contextmanager
def expect_warning(category, message):
    """Silence a known third-party warning and raise if it stops firing.
 
    Notifies us to drop the workaround once the upstream library is fixed.
    """
    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        yield
    seen = False
    for w in caught:
        if issubclass(w.category, category) and message in str(w.message):
            seen = True
        else:
            warnings.warn_explicit(w.message, w.category, w.filename, w.lineno)
    if not seen:
        raise RuntimeError(
            f"Expected {category.__name__} containing {message!r} was not "
            "emitted. This tutorial's workaround can probably be removed."
        )
 
 
# -----------------------------------------------------------------------------
# Step 1: Create and train a simple Keras model (via embedded Python)
# -----------------------------------------------------------------------------
 
input = Input(shape=(4,), batch_size=2)
x = Dense(32)(input)
x = Activation("relu")(x)
x = Dense(16, activation="relu")(x)
x = Dense(8, activation="relu")(x)
x = Dense(2)(x)
output = Softmax()(x)
model = Model(inputs=input, outputs=output)
 
randomGenerator = np.random.RandomState(0)
x_train = randomGenerator.rand(4, 4)
y_train = randomGenerator.rand(4, 2)
 
model.compile(loss="mse", optimizer="adam")
model.fit(x_train, y_train, epochs=3, batch_size=2)
 
# Keras' internal ``np.array(x)`` (TensorFlow backend) does not yet implement
# the NumPy 2.0 ``__array__(copy=...)`` signature, so saving the model emits a
# DeprecationWarning that we cannot fix from user code.
if tuple(int(p) for p in np.__version__.split(".")[:2]) >= (2, 0):
    ctx = expect_warning(DeprecationWarning, "__array__ implementation doesn't accept a copy keyword")
else:
    ctx = contextlib.nullcontext()
 
with ctx:
    model.save("KerasModel.keras")
 
model.summary()
 
# -----------------------------------------------------------------------------
# Step 2: Use TMVA::SOFIE to parse the ONNX model
# -----------------------------------------------------------------------------
 
# Parse the ONNX model
 
model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse("KerasModel.keras")
 
# Generate inference code
model.Generate()
model.OutputGenerated()
# print generated code
print("\n**************************************************")
print("           Generated code")
print("**************************************************\n")
model.PrintGenerated()
print("**************************************************\n\n\n")
 
# Compile the generated code
ROOT.gInterpreter.Declare('#include "KerasModel.hxx"')
 
 
# -----------------------------------------------------------------------------
# Step 3: Run inference
# -----------------------------------------------------------------------------
 
# instantiate SOFIE session class
session = ROOT.TMVA_SOFIE_KerasModel.Session()
 
# Input tensor (same shape as training input)
x = np.array([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]], dtype=np.float32)
 
# Run inference
y = session.infer(x)
 
print("Inference output:", y)

Epoch 1/3
 
␛[1m1/2␛[0m ␛[32m━━━━━━━━━━␛[0m␛[37m━━━━━━━━━━␛[0m ␛[1m0s␛[0m 539ms/step - loss: 0.1046␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈
␛[1m2/2␛[0m ␛[32m━━━━━━━━━━━━━━━━━━━━␛[0m␛[37m␛[0m ␛[1m1s␛[0m 11ms/step - loss: 0.1197 
Epoch 2/3
 
␛[1m1/2␛[0m ␛[32m━━━━━━━━━━␛[0m␛[37m━━━━━━━━━━␛[0m ␛[1m0s␛[0m 12ms/step - loss: 0.1368␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈
␛[1m2/2␛[0m ␛[32m━━━━━━━━━━━━━━━━━━━━␛[0m␛[37m␛[0m ␛[1m0s␛[0m 10ms/step - loss: 0.1174
Epoch 3/3
 
␛[1m1/2␛[0m ␛[32m━━━━━━━━━━␛[0m␛[37m━━━━━━━━━━␛[0m ␛[1m0s␛[0m 11ms/step - loss: 0.1351␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈␈
␛[1m2/2␛[0m ␛[32m━━━━━━━━━━━━━━━━━━━━␛[0m␛[37m␛[0m ␛[1m0s␛[0m 10ms/step - loss: 0.1151
Model: "functional"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Layer (type)                    ┃ Output Shape           ┃       Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ input_layer (InputLayer)        │ (2, 4)                 │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense (Dense)                   │ (2, 32)                │           160 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ activation (Activation)         │ (2, 32)                │             0 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_1 (Dense)                 │ (2, 16)                │           528 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_2 (Dense)                 │ (2, 8)                 │           136 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ dense_3 (Dense)                 │ (2, 2)                 │            18 │
├─────────────────────────────────┼────────────────────────┼───────────────┤
│ softmax (Softmax)               │ (2, 2)                 │             0 │
└─────────────────────────────────┴────────────────────────┴───────────────┘
 Total params: 2,528 (9.88 KB)
 Trainable params: 842 (3.29 KB)
 Non-trainable params: 0 (0.00 B)
 Optimizer params: 1,686 (6.59 KB)
PyKeras: parsing model  KerasModel.keras
 
**************************************************
           Generated code
**************************************************
 
//Code generated automatically by TMVA for Inference of Model file [KerasModel.keras] at [Mon Aug  3 13:26:29 202] 
 
#ifndef ROOT_TMVA_SOFIE_KERASMODEL
#define ROOT_TMVA_SOFIE_KERASMODEL
 
#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_KerasModel{
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 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_layer, float *tensor_keras_tensor_18 );
struct Session {
// initialized (weights and constant) tensors
std::vector<float> fTensor_dense_3kernel = std::vector<float>(16);
float * tensor_dense_3kernel = fTensor_dense_3kernel.data();
std::vector<float> fTensor_dense_2bias = std::vector<float>(8);
float * tensor_dense_2bias = fTensor_dense_2bias.data();
std::vector<float> fTensor_dense_3bias = std::vector<float>(2);
float * tensor_dense_3bias = fTensor_dense_3bias.data();
std::vector<float> fTensor_dense_2kernel = std::vector<float>(128);
float * tensor_dense_2kernel = fTensor_dense_2kernel.data();
std::vector<float> fTensor_dense_1bias = std::vector<float>(16);
float * tensor_dense_1bias = fTensor_dense_1bias.data();
std::vector<float> fTensor_dense_1kernel = std::vector<float>(512);
float * tensor_dense_1kernel = fTensor_dense_1kernel.data();
std::vector<float> fTensor_densebias = std::vector<float>(32);
float * tensor_densebias = fTensor_densebias.data();
std::vector<float> fTensor_densekernel = std::vector<float>(128);
float * tensor_densekernel = fTensor_densekernel.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 keras_tensor_8 with size 256 bytes
float* tensor_keras_tensor_8 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 0);
 
 // Allocating memory for intermediate tensor keras_tensor_10 with size 256 bytes
float* tensor_keras_tensor_10 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 256);
 
 // Allocating memory for intermediate tensor dense_1Dense with size 128 bytes
float* tensor_dense_1Dense = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 128);
 
 // Allocating memory for intermediate tensor keras_tensor_12 with size 128 bytes
float* tensor_keras_tensor_12 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 0);
 
 // Allocating memory for intermediate tensor dense_2Dense with size 64 bytes
float* tensor_dense_2Dense = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 448);
 
 // Allocating memory for intermediate tensor keras_tensor_14 with size 64 bytes
float* tensor_keras_tensor_14 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 384);
 
 // Allocating memory for intermediate tensor keras_tensor_16 with size 16 bytes
float* tensor_keras_tensor_16 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 368);
 
 // Allocating memory for intermediate tensor keras_tensor_18 with size 16 bytes
float* tensor_keras_tensor_18 = reinterpret_cast<float*>(fIntermediateMemoryPool.data() + 352);
 
 
Session(std::string filename ="KerasModel.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_dense_3kernel, "tensor_dense_3kernel", 16);
   ReadTensorFromStream(f, tensor_dense_2bias, "tensor_dense_2bias", 8);
   ReadTensorFromStream(f, tensor_dense_3bias, "tensor_dense_3bias", 2);
   ReadTensorFromStream(f, tensor_dense_2kernel, "tensor_dense_2kernel", 128);
   ReadTensorFromStream(f, tensor_dense_1bias, "tensor_dense_1bias", 16);
   ReadTensorFromStream(f, tensor_dense_1kernel, "tensor_dense_1kernel", 512);
   ReadTensorFromStream(f, tensor_densebias, "tensor_densebias", 32);
   ReadTensorFromStream(f, tensor_densekernel, "tensor_densekernel", 128);
   f.close();
 
}
 
 
 
std::vector<float> infer(float const* tensor_input_layer){
   std::vector<float > output_tensor_keras_tensor_18(4);
   doInfer(*this, tensor_input_layer, output_tensor_keras_tensor_18.data() );
   return {output_tensor_keras_tensor_18};
}
};   // end of Session
 
 
// Input tensor dimensions
constexpr std::array<SingleDim, 2> dim_input_layer{SingleDim{2}, SingleDim{4}};
 
constexpr std::array<TensorDims, 1> inputTensorDims{
   makeDims(dim_input_layer)
};
 
constexpr bool hasDynamicInputTensors{false};
 
 
// Output tensor dimensions
constexpr std::array<SingleDim, 2> dim_keras_tensor_18{SingleDim{2}, SingleDim{2}};
 
constexpr std::array<TensorDims, 1> outputTensorDims{
   makeDims(dim_keras_tensor_18)
};
 
constexpr bool hasDynamicOutputTensors{false};
 
inline void doInfer(Session const &session, float const* tensor_input_layer, float *tensor_keras_tensor_18 ) {
 
    auto &tensor_dense_1Dense = session.tensor_dense_1Dense;
    auto &tensor_dense_1bias = session.tensor_dense_1bias;
    auto &tensor_dense_1kernel = session.tensor_dense_1kernel;
    auto &tensor_dense_2Dense = session.tensor_dense_2Dense;
    auto &tensor_dense_2bias = session.tensor_dense_2bias;
    auto &tensor_dense_2kernel = session.tensor_dense_2kernel;
    auto &tensor_dense_3bias = session.tensor_dense_3bias;
    auto &tensor_dense_3kernel = session.tensor_dense_3kernel;
    auto &tensor_densebias = session.tensor_densebias;
    auto &tensor_densekernel = session.tensor_densekernel;
    auto &tensor_keras_tensor_10 = session.tensor_keras_tensor_10;
    auto &tensor_keras_tensor_12 = session.tensor_keras_tensor_12;
    auto &tensor_keras_tensor_14 = session.tensor_keras_tensor_14;
    auto &tensor_keras_tensor_16 = session.tensor_keras_tensor_16;
    auto &tensor_keras_tensor_8 = session.tensor_keras_tensor_8;
 
 
//--------- Gemm op_0 { 2 , 4 } * { 4 , 32 } -> { 2 , 32 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 32 * j;
      Copy(tensor_keras_tensor_8 + y_index, tensor_densebias, 32);
   }
   Gemm_Call(tensor_keras_tensor_8, false, false, 32, 2, 4, 1, tensor_densekernel, tensor_input_layer, 1,nullptr);
 
//------ RELU
   for (int id = 0; id < 64 ; id++){
      tensor_keras_tensor_10[id] = ((tensor_keras_tensor_8[id] > 0 )? tensor_keras_tensor_8[id] : 0);
   }
 
//--------- Gemm op_2 { 2 , 32 } * { 32 , 16 } -> { 2 , 16 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 16 * j;
      Copy(tensor_dense_1Dense + y_index, tensor_dense_1bias, 16);
   }
   Gemm_Call(tensor_dense_1Dense, false, false, 16, 2, 32, 1, tensor_dense_1kernel, tensor_keras_tensor_10, 1,nullptr);
 
//------ RELU
   for (int id = 0; id < 32 ; id++){
      tensor_keras_tensor_12[id] = ((tensor_dense_1Dense[id] > 0 )? tensor_dense_1Dense[id] : 0);
   }
 
//--------- Gemm op_4 { 2 , 16 } * { 16 , 8 } -> { 2 , 8 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 8 * j;
      Copy(tensor_dense_2Dense + y_index, tensor_dense_2bias, 8);
   }
   Gemm_Call(tensor_dense_2Dense, false, false, 8, 2, 16, 1, tensor_dense_2kernel, tensor_keras_tensor_12, 1,nullptr);
 
//------ RELU
   for (int id = 0; id < 16 ; id++){
      tensor_keras_tensor_14[id] = ((tensor_dense_2Dense[id] > 0 )? tensor_dense_2Dense[id] : 0);
   }
 
//--------- Gemm op_6 { 2 , 8 } * { 8 , 2 } -> { 2 , 2 }
   for (size_t j = 0; j < 2; j++) { 
      size_t y_index = 2 * j;
      Copy(tensor_keras_tensor_16 + y_index, tensor_dense_3bias, 2);
   }
   Gemm_Call(tensor_keras_tensor_16, false, false, 2, 2, 8, 1, tensor_dense_3kernel, tensor_keras_tensor_14, 1,nullptr);
///------- Softmax op_7 ---> { 2 , 2 }
 
   //-----  softmax axis is last one - 1
   for (int i = 0; i < 2; ++i) {
      size_t offset = i * 2;
      float const * x_ptr = &tensor_keras_tensor_16[offset];
      float * y_ptr = &tensor_keras_tensor_18[offset];
      float vmax = x_ptr[0];
      for (int j = 1; j < 2; ++j) {
         if (x_ptr[j] > vmax) vmax = x_ptr[j];
      }
      float sum = 0.0;
      for (int j = 0; j < 2; ++j) {
         y_ptr[j] = std::exp(x_ptr[j] - vmax);
         sum += y_ptr[j];
      }
      float inv_sum = 1.0f / sum;
      for (int j = 0; j < 2; ++j) {
         y_ptr[j] *= inv_sum;
      }
   }
 
}
} //TMVA_SOFIE_KerasModel
 
 
namespace clad {
namespace custom_derivatives {
namespace TMVA_SOFIE_KerasModel {
using ::TMVA_SOFIE_KerasModel::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;
   }
}
} // namespace TMVA_SOFIE_KerasModel
} // namespace custom_derivatives
} // namespace clad
 
#endif  // ROOT_TMVA_SOFIE_KERASMODEL
**************************************************
 
 
 
Inference output: { 0.520199f, 0.479801f, 0.492638f, 0.507362f }

Author

Sanjiban Sengupta and Lorenzo Moneta

Definition in file TMVA_SOFIE_Keras.py.