TMVA_SOFIE_PyTorch_HiggsModel.py

Go to the documentation of this file.

### \file
### \ingroup tutorial_ml
### \notebook -nodraw
### This macro trains a simple deep neural network on the Higgs dataset with
### PyTorch, exports the model to ONNX and runs the SOFIE parser on it to
### generate and compile C++ inference code.
###
### The trained model is saved as HiggsModel.onnx and is used as input by
### other SOFIE tutorials (e.g. TMVA_SOFIE_RDataFrame.C), so this macro needs
### to be run before them.
###
### \macro_code
### \macro_output
 
import inspect
 
import numpy as np
import ROOT
import torch
import torch.nn as nn
 
 
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 CreateModel(nlayers=4, nunits=64):
    layers = []
    ninputs = 7
    for i in range(1, 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)
 
    torch.onnx.export(model, dummy_x, modelFile, **kwargs)
 
    print("model exported to ONNX as", modelFile)
    return modelFile
 
 
def GenerateCode(modelFile="model.onnx"):
 
    # parse the input ONNX model into an RModel object
    parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX()
    model = parser.Parse(modelFile)
 
    # Generating inference code
    model.Generate()
    model.OutputGenerated()
 
    modelName = modelFile.replace(".onnx", "")
    return modelName
 
 
###################################################################
## Step 1 : Create and train the model, export it to ONNX
###################################################################
 
x_train, y_train, x_test, y_test = PrepareData()
# validate the exported model on the first test events
x_check = x_test[:10]
 
# create dense model with 3 layers of 64 units and train it
model = CreateModel(3, 64)
TrainModel(model, x_train, y_train)
modelFile = ExportModel(model, "HiggsModel")
 
# evaluate the trained model on the validation inputs, for comparison with SOFIE
with torch.no_grad():
    ytorch = model(torch.from_numpy(x_check)).numpy()
 
###################################################################
## Step 2 : Parse model and generate inference code with SOFIE
###################################################################
 
modelName = GenerateCode(modelFile)
modelHeaderFile = modelName + ".hxx"
 
###################################################################
## Step 3 : Compile the generated C++ model code
###################################################################
 
ROOT.gInterpreter.Declare('#include "' + modelHeaderFile + '"')
 
###################################################################
## Step 4: Evaluate the model
###################################################################
 
# get first the SOFIE session namespace
sofie = getattr(ROOT, "TMVA_SOFIE_" + modelName)
session = sofie.Session()
 
for i in range(x_check.shape[0]):
    y = session.infer(x_check[i])
    print("input to model is ", x_check[i], "\n\t -> output using SOFIE = ", y[0], " using PyTorch = ", ytorch[i, 0])
    if abs(y[0] - ytorch[i, 0]) > 0.01:
        raise RuntimeError("ERROR: Result is different between SOFIE and PyTorch")
 
print("OK")