doc/PR/TMVA__SOFIE__ONNX_8py.html

import os

import sys

import subprocess


import numpy as np

import ROOT


# The PyTorch export, as a small standalone script run in its own process.

# It takes the model name as its only argument and writes <modelName>.onnx.

EXPORT_SCRIPT = r"""

import sys

import inspect

import warnings

import contextlib


import torch

import torch.nn as nn


modelName = sys.argv[1]


@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."

        )


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 (dynamo path) pickles its export program through

        # copyreg, which still references the deprecated LeafSpec. The warning

        # is emitted from inside PyTorch and cannot be avoided from user code.

        with expect_warning(FutureWarning, "isinstance(treespec, LeafSpec)"):

            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__)

        # leave no .onnx behind: which the parent process treats as a RuntimeError

        sys.exit()


CreateAndTrainModel(modelName)

"""


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

##          (done in a separate process to avoid the protobuf clash)

###################################################################


# use an arbitrary modelName

modelName = "LinearModel"

modelFile = modelName + ".onnx"


subprocess.run([sys.executable, "-c", EXPORT_SCRIPT, modelName])

if not os.path.exists(modelFile):

    raise RuntimeError("ONNX model could not be exported")


###################################################################

## 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")