TMVA_SOFIE_Models.py

Go to the documentation of this file.

### \file
### \ingroup tutorial_ml
### \notebook -nodraw
### 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.
###
### \macro_code
### \macro_output
### \author Lorenzo Moneta
 
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()