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 [Mon Aug  3 13:26:51 202] 
 
#ifndef ROOT_TMVA_SOFIE_PYTORCHMODEL
#define ROOT_TMVA_SOFIE_PYTORCHMODEL
 
#include <cassert>
#include <algorithm>
#include <iomanip>
#include <cstring>
#include <vector>
#include <cstdint>
#include <limits>
#include <cmath>
#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_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
 
// --- 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_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");
   }
   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
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;
      Copy(tensor_input0 + y_index, tensor_0bias, 16);
   }
   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;
      Copy(tensor_input2 + y_index, tensor_2bias, 8);
   }
   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
 
 
namespace clad {
namespace custom_derivatives {
namespace TMVA_SOFIE_PyTorchModel {
using ::TMVA_SOFIE_PyTorchModel::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_PyTorchModel
} // namespace custom_derivatives
} // namespace clad
 
#endif  // ROOT_TMVA_SOFIE_PYTORCHMODEL

Author

Sanjiban Sengupta

Definition in file TMVA_SOFIE_PyTorch.py.