TMVA_SOFIE_ONNX.py File Reference

Detailed Description

This macro provides a simple example for:

• creating a model with Pytorch and export to ONNX
• parsing the ONNX file with SOFIE and generate C++ code
• compiling the model using ROOT Cling
• run the code and optionally compare with ONNXRuntime

 
 
import torch
import torch.nn as nn
import ROOT
import numpy as np
import inspect
 
def CreateAndTrainModel(modelName):
 
   model = nn.Sequential(
           nn.Linear(32,16),
           nn.ReLU(),
           nn.Linear(16,8),
           nn.ReLU(),
           nn.Linear(8,2),
           nn.Softmax(dim=1)
           )
 
   criterion = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(),lr=0.01)
 
 
   #train model with the random data
   for i in range(500):
      x=torch.randn(2,32)
      y=torch.randn(2,2)
      y_pred = model(x)
      loss = criterion(y_pred,y)
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
 
   #*******************************************************
   ##  EXPORT to ONNX
   #
   #  need to evaluate the model before exporting to ONNX
   #  and to provide a dummy input tensor to set the input model shape
   model.eval()
 
   modelFile = modelName + ".onnx"
   dummy_x = torch.randn(1,32)
   model(dummy_x)
 
   #check for torch.onnx.export parameters
   def filtered_kwargs(func, **candidate_kwargs):
    sig = inspect.signature(func)
    return {
        k: v for k, v in candidate_kwargs.items()
        if k in sig.parameters
   }
   kwargs = filtered_kwargs(
      torch.onnx.export,
      input_names=["input"],
      output_names=["output"],
      external_data=False,  # may not exist
      dynamo=True           # may not exist
   )
   print("calling torch.onnx.export with parameters",kwargs)
 
   try:
      torch.onnx.export(model, dummy_x, modelFile, **kwargs)
      print("model exported to ONNX as",modelFile)
      return modelFile
   except TypeError:
      print("Cannot export model from pytorch to ONNX - with version ",torch.__version__)
      print("Skip tutorial execution")
      exit()
 
 
def ParseModel(modelFile, verbose=False):
 
   parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX()
   model = parser.Parse(modelFile,verbose)
   #
   #print model weights
   if (verbose):
      model.PrintInitializedTensors()
      data = model.GetTensorData['float']('0weight')
      print("0weight",data)
      data = model.GetTensorData['float']('2weight')
      print("2weight",data)
 
   # Generating inference code
   model.Generate();
   #generate header file (and .dat file) with modelName+.hxx
   model.OutputGenerated();
   if (verbose) :
       model.PrintGenerated()
 
   modelCode = modelFile.replace(".onnx",".hxx")
   print("Generated model header file ",modelCode)
   return modelCode
 
###################################################################
## Step 1 : Create and Train model
###################################################################
 
#use an arbitrary modelName
modelName = "LinearModel"
modelFile = CreateAndTrainModel(modelName)
 
 
###################################################################
## Step 2 : Parse model and generate inference code with SOFIE
###################################################################
 
modelCode = ParseModel(modelFile, False)
 
###################################################################
## Step 3 : Compile the generated C++ model code
###################################################################
 
ROOT.gInterpreter.Declare('#include "' + modelCode + '"')
 
###################################################################
## Step 4: Evaluate the model
###################################################################
 
#get first the SOFIE session namespace
sofie = getattr(ROOT, 'TMVA_SOFIE_' + modelName)
session = sofie.Session()
 
x = np.random.normal(0,1,(1,32)).astype(np.float32)
print("\n************************************************************")
print("Running inference with SOFIE ")
print("\ninput to model is ",x)
y = session.infer(x)
# output shape is (1,2)
y_sofie = np.asarray(y.data())
print("-> output using SOFIE = ", y_sofie)
 
#check inference with onnx
try:
   import onnxruntime as ort
    # Load model
   print("Running inference with ONNXRuntime ")
   ort_session = ort.InferenceSession(modelFile)
 
   # Run inference
   outputs = ort_session.run(None, {"input": x})
   y_ort = outputs[0]
   print("-> output using ORT =", y_ort)
 
   testFailed =  abs(y_sofie-y_ort) > 0.01
   if (np.any(testFailed)):
      raiseError('Result is different between SOFIE and ONNXRT')
   else :
      print("OK")
 
except ImportError:
   print("Missing ONNXRuntime: skipping comparison test")

calling torch.onnx.export with parameters {'input_names': ['input'], 'output_names': ['output'], 'external_data': False, 'dynamo': True}
[torch.onnx] Obtain model graph for `Sequential([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `Sequential([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decomposition...
[torch.onnx] Run decomposition... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
model exported to ONNX as LinearModel.onnx
Generated model header file  LinearModel.hxx
 
************************************************************
Running inference with SOFIE 
 
input to model is  [[-6.01221144e-01  1.19053714e-01 -2.45943832e+00  2.12796196e-01
   3.59251350e-01  1.69989944e+00 -1.45543015e+00 -2.42614836e-01
   2.40988610e-03 -5.08478820e-01  3.93389344e+00  1.12916040e+00
  -4.54266518e-01 -6.74715400e-01  5.57646573e-01  1.61057740e-01
   4.01047081e-01 -2.17865658e+00 -4.93074208e-01  4.91779983e-01
   6.03726096e-02  1.46321520e-01 -8.50548074e-02  3.65753442e-01
  -1.35761845e+00  7.16157615e-01  2.67857730e-01  1.38448700e-01
   1.31619394e+00  3.28651786e-01  2.88355619e-01  1.74426639e+00]]
-> output using SOFIE =  [0.41176823 0.5882318 ]
Missing ONNXRuntime: skipping comparison test

Author

Lorenzo Moneta

Definition in file TMVA_SOFIE_ONNX.py.