TMVA_SOFIE_Models.py File Reference

Detailed Description

Example of inference with SOFIE using a set of models trained with PyTorch and exported to ONNX.

This tutorial shows how to store several models in a single header file and the weights in a ROOT binary file. The models are then evaluated using the RDataFrame

The PyTorch export and ROOT's SOFIE parser are both linked against protobuf, but usually against different versions, so loading them in the same process leads to a symbol clash. We therefore run the PyTorch training and ONNX export in a separate Python process and only use ROOT before and afterwards.

 
import os
import subprocess
import sys
 
import numpy as np
import ROOT
 
## generate and train PyTorch models with different architectures
 
# The PyTorch training and ONNX export, as a small standalone script run in its
# own process. It takes as arguments the .npz file with the training data and
# the names of the models to train, and writes a <modelName>.onnx file for each
# of them.
TRAIN_SCRIPT = r"""
import sys
import inspect
import warnings
import contextlib
 
import numpy as np
import torch
import torch.nn as nn
 
dataFile = sys.argv[1]
modelNames = sys.argv[2:]
 
 
@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 CreateModel(nlayers=4, nunits=64):
    layers = []
    ninputs = 7
    for i in range(nlayers):
        layers += [nn.Linear(ninputs, nunits), nn.ReLU()]
        ninputs = nunits
    layers += [nn.Linear(ninputs, 1), nn.Sigmoid()]
    model = nn.Sequential(*layers)
    print(model)
    return model
 
 
def TrainModel(model, x, y, epochs=5, batch_size=50):
    x = torch.from_numpy(x)
    y = torch.from_numpy(y)
    criterion = nn.BCELoss()
    optimizer = torch.optim.Adam(model.parameters())
    nbatches = x.shape[0] // batch_size
    for epoch in range(epochs):
        perm = torch.randperm(x.shape[0])
        running_loss = 0.0
        for i in range(nbatches):
            idx = perm[i * batch_size : (i + 1) * batch_size]
            optimizer.zero_grad()
            loss = criterion(model(x[idx]), y[idx])
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
        print(f"Epoch {epoch + 1}/{epochs} - average loss: {running_loss / nbatches:.4f}")
 
 
def ExportModel(model, modelName):
    # need to evaluate the model before exporting to ONNX
    # and to provide a dummy input tensor to set the input model shape
    # (the batch size is fixed to 1 for the SOFIE inference)
    model.eval()
 
    modelFile = modelName + ".onnx"
    dummy_x = torch.randn(1, 7)
    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)
    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()
 
 
data = np.load(dataFile)
for modelName in modelNames:
    model = CreateModel(4, 64)
    TrainModel(model, data["x_train"], data["y_train"])
    ExportModel(model, modelName)
"""
 
 
def PrepareData():
    # get the input data
    inputFile = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/Higgs_data.root"
 
    df1 = ROOT.RDataFrame("sig_tree", inputFile)
    sigData = df1.AsNumpy(columns=["m_jj", "m_jjj", "m_lv", "m_jlv", "m_bb", "m_wbb", "m_wwbb"])
    # print(sigData)
 
    # stack all the 7 numpy array in a single array (nevents x nvars)
    xsig = np.column_stack(list(sigData.values()))
    data_sig_size = xsig.shape[0]
    print("size of data", data_sig_size)
 
    # make SOFIE inference on background data
    df2 = ROOT.RDataFrame("bkg_tree", inputFile)
    bkgData = df2.AsNumpy(columns=["m_jj", "m_jjj", "m_lv", "m_jlv", "m_bb", "m_wbb", "m_wwbb"])
    xbkg = np.column_stack(list(bkgData.values()))
    data_bkg_size = xbkg.shape[0]
 
    ysig = np.ones(data_sig_size)
    ybkg = np.zeros(data_bkg_size)
    inputs_data = np.concatenate((xsig, xbkg), axis=0).astype(np.float32)
    inputs_targets = np.concatenate((ysig, ybkg), axis=0).astype(np.float32)
 
    # split data in training and test data
    rng = np.random.default_rng(1234)
    idx = rng.permutation(inputs_data.shape[0])
    ntrain = inputs_data.shape[0] // 2
 
    x_train = inputs_data[idx[:ntrain]]
    y_train = inputs_targets[idx[:ntrain]].reshape(-1, 1)
    x_test = inputs_data[idx[ntrain:]]
    y_test = inputs_targets[idx[ntrain:]].reshape(-1, 1)
 
    return x_train, y_train, x_test, y_test
 
 
def TrainModels(x_train, y_train, modelNames):
    # train the models with PyTorch and export them to ONNX
    # (done in a separate process to avoid the protobuf clash, see above)
    dataFile = "Higgs_Model_train_data.npz"
    np.savez(dataFile, x_train=x_train, y_train=y_train)
 
    subprocess.run([sys.executable, "-c", TRAIN_SCRIPT, dataFile] + modelNames, check=True)
    os.remove(dataFile)
 
    modelFiles = [name + ".onnx" for name in modelNames]
    for modelFile in modelFiles:
        if not os.path.exists(modelFile):
            raise RuntimeError("ONNX model " + modelFile + " could not be exported")
    return modelFiles
 
 
### run the models
 
x_train, y_train, x_test, y_test = PrepareData()
 
## create models and train them
 
# All three models use the same small architecture (4 layers of 64 units) to
# keep the tutorial runtime under control, whatever their names suggest.
modelNames = ["Higgs_Model_4L_50", "Higgs_Model_4L_200", "Higgs_Model_2L_500"]
model1, model2, model3 = TrainModels(x_train, y_train, modelNames)
 
# evaluate with SOFIE the 3 trained models
 
 
def GenerateModelCode(modelFile, generatedHeaderFile):
    parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX()
    model = parser.Parse(modelFile)
 
    print("Generating inference code for the ONNX model from ", modelFile, "in the header ", generatedHeaderFile)
    # Generating inference code using a ROOT binary file
    model.Generate(ROOT.TMVA.Experimental.SOFIE.Options.kRootBinaryWeightFile)
    # add option to append to the same file the generated headers (pass True for append flag)
    model.OutputGenerated(generatedHeaderFile, True)
    # model.PrintGenerated()
    return generatedHeaderFile
 
 
generatedHeaderFile = "Higgs_Model.hxx"
# need to remove existing header file since we are appending on same one
if os.path.exists(generatedHeaderFile):
    print("removing existing file", generatedHeaderFile)
    os.remove(generatedHeaderFile)
 
weightFile = "Higgs_Model.root"
if os.path.exists(weightFile):
    print("removing existing file", weightFile)
    os.remove(weightFile)
 
GenerateModelCode(model1, generatedHeaderFile)
GenerateModelCode(model2, generatedHeaderFile)
GenerateModelCode(model3, generatedHeaderFile)
 
# compile the generated code
 
ROOT.gInterpreter.Declare('#include "' + generatedHeaderFile + '"')
 
 
# run the inference on the test data
session1 = ROOT.TMVA_SOFIE_Higgs_Model_4L_50.Session("Higgs_Model.root")
session2 = ROOT.TMVA_SOFIE_Higgs_Model_4L_200.Session("Higgs_Model.root")
session3 = ROOT.TMVA_SOFIE_Higgs_Model_2L_500.Session("Higgs_Model.root")
 
hs1 = ROOT.TH1D("hs1", "Signal result 4L 50", 100, 0, 1)
hs2 = ROOT.TH1D("hs2", "Signal result 4L 200", 100, 0, 1)
hs3 = ROOT.TH1D("hs3", "Signal result 2L 500", 100, 0, 1)
 
hb1 = ROOT.TH1D("hb1", "Background result 4L 50", 100, 0, 1)
hb2 = ROOT.TH1D("hb2", "Background result 4L 200", 100, 0, 1)
hb3 = ROOT.TH1D("hb3", "Background result 2L 500", 100, 0, 1)
 
 
def EvalModel(session, x):
    result = session.infer(x)
    return result[0]
 
 
for i in range(0, x_test.shape[0]):
    result1 = EvalModel(session1, x_test[i, :])
    result2 = EvalModel(session2, x_test[i, :])
    result3 = EvalModel(session3, x_test[i, :])
    if y_test[i] == 1:
        hs1.Fill(result1)
        hs2.Fill(result2)
        hs3.Fill(result3)
    else:
        hb1.Fill(result1)
        hb2.Fill(result2)
        hb3.Fill(result3)
 
 
def PlotHistos(hs, hb):
    hs.SetLineColor("kRed")
    hb.SetLineColor("kBlue")
    hs.Draw()
    hb.Draw("same")
 
 
c1 = ROOT.TCanvas()
c1.Divide(1, 3)
c1.cd(1)
PlotHistos(hs1, hb1)
c1.cd(2)
PlotHistos(hs2, hb2)
c1.cd(3)
PlotHistos(hs3, hb3)
c1.Draw()
 
## draw also ROC curves
 
 
def GetContent(h):
    n = h.GetNbinsX()
    x = ROOT.std.vector["float"](n)
    w = ROOT.std.vector["float"](n)
    for i in range(0, n):
        x[i] = h.GetBinCenter(i + 1)
        w[i] = h.GetBinContent(i + 1)
    return x, w
 
 
def MakeROCCurve(hs, hb):
    xs, ws = GetContent(hs)
    xb, wb = GetContent(hb)
    roc = ROOT.TMVA.ROCCurve(xs, xb, ws, wb)
    print("ROC integral for ", hs.GetName(), roc.GetROCIntegral())
    curve = roc.GetROCCurve()
    curve.SetName(hs.GetName())
    return roc, curve
 
 
c2 = ROOT.TCanvas()
 
r1, curve1 = MakeROCCurve(hs1, hb1)
curve1.SetLineColor("kRed")
curve1.Draw("AC")
 
r2, curve2 = MakeROCCurve(hs2, hb2)
curve2.SetLineColor("kBlue")
curve2.Draw("C")
 
r3, curve3 = MakeROCCurve(hs3, hb3)
curve3.SetLineColor("kGreen")
curve3.Draw("C")
 
c2.Draw()

size of data 10000
Sequential(
  (0): Linear(in_features=7, out_features=64, bias=True)
  (1): ReLU()
  (2): Linear(in_features=64, out_features=64, bias=True)
  (3): ReLU()
  (4): Linear(in_features=64, out_features=64, bias=True)
  (5): ReLU()
  (6): Linear(in_features=64, out_features=64, bias=True)
  (7): ReLU()
  (8): Linear(in_features=64, out_features=1, bias=True)
  (9): Sigmoid()
)
Epoch 1/5 - average loss: 0.6688
Epoch 2/5 - average loss: 0.6392
Epoch 3/5 - average loss: 0.6257
Epoch 4/5 - average loss: 0.6142
Epoch 5/5 - average loss: 0.6101
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 Higgs_Model_4L_50.onnx
Sequential(
  (0): Linear(in_features=7, out_features=64, bias=True)
  (1): ReLU()
  (2): Linear(in_features=64, out_features=64, bias=True)
  (3): ReLU()
  (4): Linear(in_features=64, out_features=64, bias=True)
  (5): ReLU()
  (6): Linear(in_features=64, out_features=64, bias=True)
  (7): ReLU()
  (8): Linear(in_features=64, out_features=1, bias=True)
  (9): Sigmoid()
)
Epoch 1/5 - average loss: 0.6743
Epoch 2/5 - average loss: 0.6507
Epoch 3/5 - average loss: 0.6357
Epoch 4/5 - average loss: 0.6267
Epoch 5/5 - average loss: 0.6204
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 Higgs_Model_4L_200.onnx
Sequential(
  (0): Linear(in_features=7, out_features=64, bias=True)
  (1): ReLU()
  (2): Linear(in_features=64, out_features=64, bias=True)
  (3): ReLU()
  (4): Linear(in_features=64, out_features=64, bias=True)
  (5): ReLU()
  (6): Linear(in_features=64, out_features=64, bias=True)
  (7): ReLU()
  (8): Linear(in_features=64, out_features=1, bias=True)
  (9): Sigmoid()
)
Epoch 1/5 - average loss: 0.6680
Epoch 2/5 - average loss: 0.6389
Epoch 3/5 - average loss: 0.6237
Epoch 4/5 - average loss: 0.6171
Epoch 5/5 - average loss: 0.6108
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 Higgs_Model_2L_500.onnx
Generating inference code for the ONNX model from  Higgs_Model_4L_50.onnx in the header  Higgs_Model.hxx
Generating inference code for the ONNX model from  Higgs_Model_4L_200.onnx in the header  Higgs_Model.hxx
Generating inference code for the ONNX model from  Higgs_Model_2L_500.onnx in the header  Higgs_Model.hxx
ROC integral for  hs1 0.7230905556009574
ROC integral for  hs2 0.7247763204017561
ROC integral for  hs3 0.7310298807886174

Author

Lorenzo Moneta

Definition in file TMVA_SOFIE_Models.py.