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 inspect
 
import numpy as np
import ROOT
import torch
import torch.nn as nn
 
 
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)
 
    torch.onnx.export(model, dummy_x, modelFile, **kwargs)
 
    print("model exported to ONNX as", modelFile)
    return modelFile
 
 
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 the model, export it to ONNX
###################################################################
 
# 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):
        raise RuntimeError("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 decompositions...
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅
model exported to ONNX as LinearModel.onnx
Generated model header file  LinearModel.hxx
 
************************************************************
Running inference with SOFIE 
 
input to model is  [[ 1.367764    0.50571823  0.656371   -0.1435604   2.6970513   0.6069507
  -0.08378253 -0.08815525 -0.9546562   0.36806196  1.5668218  -0.8092264
   1.2457751   0.5912777  -1.4591385   0.6704636   0.9441589  -0.66921526
  -0.53377193 -0.37346548  0.4413543   1.0838823   0.26820326 -0.02907695
  -0.44410625  0.94629586  0.7192299  -0.26345798  0.47244576 -1.5309932
  -1.0469879  -0.23127167]]
-> output using SOFIE =  [0.4874975  0.51250255]
Missing ONNXRuntime: skipping comparison test

Author

Lorenzo Moneta

Definition in file TMVA_SOFIE_ONNX.py.