From 593322eb486b269f721436528a0b190dfbd87840 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Mon, 13 Jul 2026 09:42:12 +0200 Subject: [PATCH] [tmva][sofie] Migrate SOFIE tutorials from Keras to PyTorch and ONNX MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PyTorch has become the dominant training framework in the particle physics community, and it comes with better ONNX support than Keras. Since ONNX is the primary input format for SOFIE, the SOFIE tutorials based on the Higgs dataset are migrated from Keras to PyTorch with ONNX export. Replace TMVA_SOFIE_Keras_HiggsModel.py with TMVA_SOFIE_PyTorch_HiggsModel.py, which trains the same architecture (7 -> 64 -> 64 -> 1 with sigmoid output, binary cross-entropy, Adam, 5 epochs) with PyTorch, exports it to HiggsModel.onnx and generates the inference code with RModelParser_ONNX. The generated code is validated against the PyTorch predictions, and the trained model was checked to be roughly equivalent to the Keras one (test AUC 0.712 vs. 0.720 on an identical data split). Adapt the downstream tutorials that consume the trained Higgs model (TMVA_SOFIE_RDataFrame.C/.py, TMVA_SOFIE_RDataFrame_JIT.C, TMVA_SOFIE_RSofieReader.C, TMVA_SOFIE_Inference.py) to use HiggsModel.onnx instead of HiggsModel.keras. Also migrate TMVA_SOFIE_Models.py, which demonstrates storing several models in a single generated header with weights in one ROOT binary file: this feature is independent of the training framework. The PyTorch ONNX export and ROOT's SOFIE parser are usually linked against different protobuf versions, so the training and export run in a separate Python process, following the pattern established in TMVA_SOFIE_ONNX.py. The CMake test dependencies are updated accordingly, and the migrated tutorials are now gated on tmva-sofie plus the torch and onnx Python modules instead of keras. 🤖 Done with the help of AI. --- tmva/sofie/README.md | 2 +- tutorials/CMakeLists.txt | 28 +- .../machine_learning/TMVA_SOFIE_Inference.py | 26 +- .../TMVA_SOFIE_Keras_HiggsModel.py | 164 ----------- .../machine_learning/TMVA_SOFIE_Models.py | 175 +++++++++--- .../TMVA_SOFIE_PyTorch_HiggsModel.py | 257 ++++++++++++++++++ .../machine_learning/TMVA_SOFIE_RDataFrame.C | 10 +- .../machine_learning/TMVA_SOFIE_RDataFrame.py | 13 +- .../TMVA_SOFIE_RDataFrame_JIT.C | 11 +- .../TMVA_SOFIE_RSofieReader.C | 14 +- tutorials/machine_learning/index.md | 12 +- 11 files changed, 442 insertions(+), 270 deletions(-) delete mode 100644 tutorials/machine_learning/TMVA_SOFIE_Keras_HiggsModel.py create mode 100644 tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py diff --git a/tmva/sofie/README.md b/tmva/sofie/README.md index 79043be5c7bce..f84653a257b85 100644 --- a/tmva/sofie/README.md +++ b/tmva/sofie/README.md @@ -176,9 +176,9 @@ parser.CheckModel("example_model.ONNX"); - **Tutorials** - [TMVA_SOFIE_Inference](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_Inference.py) - [TMVA_SOFIE_Keras](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_Keras.C) - - [TMVA_SOFIE_Keras_HiggsModel](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_Keras_HiggsModel.C) - [TMVA_SOFIE_ONNX](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_ONNX.C) - [TMVA_SOFIE_PyTorch](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_PyTorch.py) + - [TMVA_SOFIE_PyTorch_HiggsModel](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py) - [TMVA_SOFIE_RDataFrame](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.C) - [TMVA_SOFIE_RDataFrame](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.py) - [TMVA_SOFIE_RDataFrame_JIT](https://github.com/root-project/root/blob/master/tutorials/machine_learning/TMVA_SOFIE_RDataFrame_JIT.C) diff --git a/tutorials/CMakeLists.txt b/tutorials/CMakeLists.txt index 3297527a27e8f..483c23a87250a 100644 --- a/tutorials/CMakeLists.txt +++ b/tutorials/CMakeLists.txt @@ -363,20 +363,20 @@ else() # These SOFIE tutorials take models trained via PyMVA-PyKeras as input if (NOT tmva-sofie OR NOT ROOT_KERAS_FOUND) list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Keras.py) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Models.py) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Keras_HiggsModel.py) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame.C) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame_JIT.C) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RSofieReader.C) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame.py) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RSofieReader.C) - list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Inference.py) endif() if (NOT tmva-sofie OR NOT ROOT_TORCH_FOUND) list(APPEND tmva_veto machine_learning/TMVA_SOFIE_PyTorch.py) endif() + # These SOFIE tutorials take the ONNX HiggsModel trained with PyTorch as input if (NOT tmva-sofie OR NOT ROOT_TORCH_FOUND OR NOT ROOT_ONNX_FOUND) list(APPEND tmva_veto machine_learning/TMVA_SOFIE_ONNX.py) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Models.py) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame.C) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame_JIT.C) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RSofieReader.C) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_RDataFrame.py) + list(APPEND tmva_veto machine_learning/TMVA_SOFIE_Inference.py) endif() #veto this tutorial since it is added directly list(APPEND tmva_veto machine_learning/TMVA_SOFIE_GNN_Parser.py) @@ -641,7 +641,7 @@ set (machine_learning-tmva004_RStandardScaler-depends tutorial-machine_learning- set (machine_learning-pytorch-ApplicationClassificationPyTorch-depends tutorial-machine_learning-pytorch-ClassificationPyTorch-py) set (machine_learning-pytorch-RegressionPyTorch-depends tutorial-machine_learning-pytorch-ApplicationClassificationPyTorch-py) set (machine_learning-pytorch-ApplicationRegressionPyTorch-depends tutorial-machine_learning-pytorch-RegressionPyTorch-py) -set (machine_learning-TMVA_SOFIE_RDataFrame-depends tutorial-machine_learning-TMVA_SOFIE_Keras_HiggsModel-py) +set (machine_learning-TMVA_SOFIE_RDataFrame-depends tutorial-machine_learning-TMVA_SOFIE_PyTorch_HiggsModel-py) set (machine_learning-TMVA_SOFIE_RDataFrame_JIT-depends tutorial-machine_learning-TMVA_SOFIE_RDataFrame) set (machine_learning-TMVA_SOFIE_RSofieReader-depends tutorial-machine_learning-TMVA_SOFIE_RDataFrame_JIT) set (machine_learning-TMVA_SOFIE_Inference-depends tutorial-machine_learning-TMVA_SOFIE_RSofieReader) @@ -834,8 +834,8 @@ if(pyroot) list(APPEND pyveto ${tmva_veto_py}) endif() - if (ROOT_KERAS_FOUND) - set (machine_learning-TMVA_SOFIE_RDataFrame-py-depends tutorial-machine_learning-TMVA_SOFIE_Keras_HiggsModel) + if (ROOT_TORCH_FOUND AND ROOT_ONNX_FOUND) + set (machine_learning-TMVA_SOFIE_RDataFrame-py-depends tutorial-machine_learning-TMVA_SOFIE_PyTorch_HiggsModel) endif() if(NOT tmva-pymva) @@ -933,9 +933,6 @@ if(pyroot) roofit/roofit/rf409_NumPyPandasToRooFit.py) file(GLOB requires_keras RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} - machine_learning/TMVA_SOFIE_Inference.py - machine_learning/TMVA_SOFIE_Models.py - machine_learning/TMVA_SOFIE_RDataFrame.py machine_learning/keras/*.py ) file(GLOB requires_torch RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} @@ -943,7 +940,9 @@ if(pyroot) machine_learning/ml_dataloader_PyTorch.py machine_learning/ml_dataloader_resampling.py machine_learning/ml_dataloader_Higgs_Classification.py + machine_learning/TMVA_SOFIE_Models.py machine_learning/TMVA_SOFIE_PyTorch.py + machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py ) file(GLOB requires_xgboost RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} machine_learning/tmva101_Training.py @@ -951,7 +950,6 @@ if(pyroot) roofit/roofit/rf618_mixture_models.py ) file(GLOB requires_sklearn RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} - machine_learning/TMVA_SOFIE_Models.py machine_learning/tmva101_Training.py # uses the xgboost sklearn plugin machine_learning/tmva102_Testing.py # requires tmva101_Training.py roofit/roofit/rf617_simulation_based_inference_multidimensional.py diff --git a/tutorials/machine_learning/TMVA_SOFIE_Inference.py b/tutorials/machine_learning/TMVA_SOFIE_Inference.py index 78edb73fa6e5e..26b3c905957f4 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_Inference.py +++ b/tutorials/machine_learning/TMVA_SOFIE_Inference.py @@ -1,10 +1,10 @@ ### \file ### \ingroup tutorial_ml ### \notebook -nodraw -### This macro provides an example of using a trained model with Keras +### This macro provides an example of using a trained model with PyTorch ### and make inference using SOFIE directly from Numpy -### This macro uses as input a Keras model generated with the -### TMVA_Higgs_Classification.C tutorial +### This macro uses as input an ONNX model generated with the +### TMVA_SOFIE_PyTorch_HiggsModel.py tutorial ### You need to run that macro before this one. ### In this case we are parsing the input file and then run the inference in the same ### macro making use of the ROOT JITing capability @@ -20,32 +20,28 @@ import ROOT # check if the input file exists -modelFile = "HiggsModel.keras" +modelFile = "HiggsModel.onnx" if not exists(modelFile): - raise FileNotFoundError("You need to run TMVA_Higgs_Classification.C to generate the Keras trained model") + raise FileNotFoundError("You need to run TMVA_SOFIE_PyTorch_HiggsModel.py to generate the ONNX trained model") -# parse the input Keras model into RModel object -model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse(modelFile) +# parse the input ONNX model into RModel object +parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX() +model = parser.Parse(modelFile) -generatedHeaderFile = modelFile.replace(".keras",".hxx") -print("Generating inference code for the Keras model from ",modelFile,"in the header ", generatedHeaderFile) +generatedHeaderFile = modelFile.replace(".onnx", ".hxx") +print("Generating inference code for the ONNX model from ", modelFile, "in the header ", generatedHeaderFile) #Generating inference code model.Generate() model.OutputGenerated(generatedHeaderFile) model.PrintGenerated() # now compile using ROOT JIT trained model -modelName = modelFile.replace(".keras","") +modelName = modelFile.replace(".onnx", "") print("compiling SOFIE model ", modelName) ROOT.gInterpreter.Declare('#include "' + generatedHeaderFile + '"') - -generatedHeaderFile = modelFile.replace(".keras",".hxx") -print("Generating inference code for the Keras model from ",modelFile,"in the header ", generatedHeaderFile) -#Generating inference - inputFileName = "Higgs_data.root" inputFile = str(ROOT.gROOT.GetTutorialDir()) + "/machine_learning/data/" + inputFileName diff --git a/tutorials/machine_learning/TMVA_SOFIE_Keras_HiggsModel.py b/tutorials/machine_learning/TMVA_SOFIE_Keras_HiggsModel.py deleted file mode 100644 index bfda3972b976c..0000000000000 --- a/tutorials/machine_learning/TMVA_SOFIE_Keras_HiggsModel.py +++ /dev/null @@ -1,164 +0,0 @@ -### \file -### \ingroup tutorial_ml -### \notebook -nodraw -### This macro run the SOFIE parser on the Keras model -### obtaining running TMVA_Higgs_Classification.C -### You need to run that macro before this one -### -### \author Lorenzo Moneta - - -import contextlib -import warnings -from os.path import exists - -import numpy as np -import ROOT -from keras import layers, models -from sklearn.model_selection import train_test_split - - -@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): - input = layers.Input(shape=(7,)) - x = input - for i in range(1, nlayers): - y = layers.Dense(nunits, activation="relu")(x) - x = y - - output = layers.Dense(1, activation="sigmoid")(x) - model = models.Model(input, output) - model.compile(loss="binary_crossentropy", optimizer="adam", weighted_metrics=["accuracy"]) - model.summary() - return model - - -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) - inputs_targets = np.concatenate((ysig, ybkg), axis=0) - - # split data in training and test data - - x_train, x_test, y_train, y_test = train_test_split(inputs_data, inputs_targets, test_size=0.50, random_state=1234) - - return x_train, y_train, x_test, y_test - - -def TrainModel(model, x, y, name): - model.fit(x, y, epochs=5, batch_size=50) - modelFile = name + ".keras" - - # Keras' internal ``np.array(x)`` (TensorFlow backend) does not yet implement - # the NumPy 2.0 ``__array__(copy=...)`` signature, so saving the model emits a - # DeprecationWarning that we cannot fix from user code. - if tuple(int(p) for p in np.__version__.split(".")[:2]) >= (2, 0): - ctx = expect_warning(DeprecationWarning, "__array__ implementation doesn't accept a copy keyword") - else: - ctx = contextlib.nullcontext() - - with ctx: - model.save(modelFile) - - return model, modelFile - - -def GenerateCode(modelFile="model.keras"): - - # check if the input file exists - if not exists(modelFile): - raise FileNotFoundError( - "INput model file not existing. You need to run TMVA_Higgs_Classification.C to generate the Keras trained model" - ) - - # parse the input Keras model into RModel object (force batch size to be 1) - model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse(modelFile) - - # Generating inference code - model.Generate() - model.OutputGenerated() - - modelName = modelFile.replace(".keras", "") - return modelName - - -################################################################### -## Step 1 : Create and Train model -################################################################### - -x_train, y_train, x_test, y_test = PrepareData() -# create dense model with 3 layers of 64 units -model = CreateModel(3, 64) -model, modelFile = TrainModel(model, x_train, y_train, "HiggsModel") - -################################################################### -## 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() - -x = np.random.normal(0, 1, 7).astype(np.float32) -y = session.infer(x) -ykeras = model(x.reshape(1, 7)).numpy() - -print("input to model is ", x, "\n\t -> output using SOFIE = ", y[0], " using Keras = ", ykeras[0]) - -if abs(y[0] - ykeras[0]) > 0.01: - raise RuntimeError("ERROR: Result is different between SOFIE and Keras") - -print("OK") diff --git a/tutorials/machine_learning/TMVA_SOFIE_Models.py b/tutorials/machine_learning/TMVA_SOFIE_Models.py index 14114a5831db7..889b5056064b4 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_Models.py +++ b/tutorials/machine_learning/TMVA_SOFIE_Models.py @@ -1,36 +1,53 @@ ### \file ### \ingroup tutorial_ml ### \notebook -nodraw -### Example of inference with SOFIE using a set of models trained with Keras. +### 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 -### First, generate the input model by running `TMVA_Higgs_Classification.C`. ### -### This tutorial parses the input model and runs the inference using ROOT's JITing capability. +### 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 contextlib import os -import warnings +import subprocess +import sys import numpy as np import ROOT -from sklearn.model_selection import train_test_split -from tensorflow.keras.layers import Dense, Input -from tensorflow.keras.models import Sequential -from tensorflow.keras.optimizers import Adam + +## 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 .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. Raise if it stops firing. + # Silence a known third-party warning and raise if it stops firing. - Notifies us to drop the workaround once the upstream library is fixed. - """ + # Notifies us to drop the workaround once the upstream library is fixed. with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") yield @@ -47,22 +64,82 @@ def expect_warning(category, message): ) -## generate and train Keras models with different architectures - - def CreateModel(nlayers=4, nunits=64): - model = Sequential() - model.add(Input(shape=(7,))) - model.add(Dense(nunits, activation="relu")) - for i in range(1, nlayers): - model.add(Dense(nunits, activation="relu")) - - model.add(Dense(1, activation="sigmoid")) - model.compile(loss="binary_crossentropy", optimizer=Adam(learning_rate=0.001), weighted_metrics=["accuracy"]) - model.summary() + 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" @@ -84,30 +161,36 @@ def PrepareData(): ysig = np.ones(data_sig_size) ybkg = np.zeros(data_bkg_size) - inputs_data = np.concatenate((xsig, xbkg), axis=0) - inputs_targets = np.concatenate((ysig, ybkg), axis=0) + 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, x_test, y_train, y_test = train_test_split(inputs_data, inputs_targets, test_size=0.50, random_state=1234) + 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 TrainModel(model, x, y, name): - model.fit(x, y, epochs=5, batch_size=50) - modelFile = name + ".keras" - # Keras' internal ``np.array(x)`` (TensorFlow backend) does not yet - # implement the NumPy 2.0 ``__array__(copy=...)`` signature, so saving - # emits a DeprecationWarning that we cannot fix from user code. - if tuple(int(p) for p in np.__version__.split(".")[:2]) >= (2, 0): - ctx = expect_warning(DeprecationWarning, "__array__ implementation doesn't accept a copy keyword") - else: - ctx = contextlib.nullcontext() +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) - with ctx: - model.save(modelFile) - return modelFile + 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 @@ -116,17 +199,19 @@ def TrainModel(model, x, y, name): ## create models and train them -model1 = TrainModel(CreateModel(4, 64), x_train, y_train, "Higgs_Model_4L_50") -model2 = TrainModel(CreateModel(4, 64), x_train, y_train, "Higgs_Model_4L_200") -model3 = TrainModel(CreateModel(4, 64), x_train, y_train, "Higgs_Model_2L_500") +# 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): - model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse(modelFile) + parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX() + model = parser.Parse(modelFile) - print("Generating inference code for the Keras model from ", modelFile, "in the header ", generatedHeaderFile) + 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) diff --git a/tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py b/tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py new file mode 100644 index 0000000000000..71629fdd4452e --- /dev/null +++ b/tutorials/machine_learning/TMVA_SOFIE_PyTorch_HiggsModel.py @@ -0,0 +1,257 @@ +### \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. +### +### 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 + +import os +import subprocess +import sys + +import numpy as np +import ROOT + +# 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 model name, and writes .onnx together with the PyTorch +# predictions for the validation inputs in _torch_output.npy. +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] +modelName = 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(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) + + 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) + +# create dense model with 3 layers of 64 units and train it +model = CreateModel(3, 64) +TrainModel(model, data["x_train"], data["y_train"]) +ExportModel(model, modelName) + +# evaluate the trained model on the validation inputs, for comparison with SOFIE +with torch.no_grad(): + y = model(torch.from_numpy(data["x_check"])).numpy() +np.save(modelName + "_torch_output.npy", y) +""" + + +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 TrainModel(x_train, y_train, x_check, name): + # train the model with PyTorch and export it to ONNX + # (done in a separate process to avoid the protobuf clash, see above) + dataFile = name + "_train_data.npz" + np.savez(dataFile, x_train=x_train, y_train=y_train, x_check=x_check) + + modelFile = name + ".onnx" + torchOutputFile = name + "_torch_output.npy" + subprocess.run([sys.executable, "-c", TRAIN_SCRIPT, dataFile, name], check=True) + os.remove(dataFile) + if not os.path.exists(modelFile) or not os.path.exists(torchOutputFile): + raise RuntimeError("ONNX model could not be exported") + + ytorch = np.load(torchOutputFile) + os.remove(torchOutputFile) + return modelFile, ytorch + + +def GenerateCode(modelFile="model.onnx"): + + # check if the input file exists + if not os.path.exists(modelFile): + raise FileNotFoundError("Input model file is missing. The PyTorch training did not produce " + modelFile) + + # 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] +modelFile, ytorch = TrainModel(x_train, y_train, x_check, "HiggsModel") + +################################################################### +## 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") diff --git a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.C b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.C index 4bc663c4019ec..433682ab858b6 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.C +++ b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.C @@ -1,16 +1,16 @@ /// \file /// \ingroup tutorial_ml /// \notebook -nodraw -/// This macro provides an example of using a trained model with Keras +/// This macro provides an example of using a trained model with PyTorch /// and make inference using SOFIE and RDataFrame -/// This macro uses as input a Keras model generated with the -/// Python tutorial TMVA_SOFIE_Keras_HiggsModel.py -/// You need to run that macro before to generate the trained Keras model +/// This macro uses as input an ONNX model generated with the +/// Python tutorial TMVA_SOFIE_PyTorch_HiggsModel.py +/// You need to run that macro before to generate the trained PyTorch model /// and also the corresponding header file with SOFIE which can then be used for inference /// /// Execute in this order: /// ``` -/// python3 TMVA_SOFIE_Keras_HiggsModel.py +/// python3 TMVA_SOFIE_PyTorch_HiggsModel.py /// root TMVA_SOFIE_RDataFrame.C /// ``` /// diff --git a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.py b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.py index af1c059fb544c..f700ce81cdf59 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.py +++ b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame.py @@ -1,8 +1,8 @@ ### \file ### \ingroup tutorial_ml ### \notebook -nodraw -### Example of inference with SOFIE and RDataFrame, of a model trained with Keras. -### First, generate the input model by running `TMVA_Higgs_Classification.C`. +### Example of inference with SOFIE and RDataFrame, of a model trained with PyTorch. +### First, generate the input ONNX model by running `TMVA_SOFIE_PyTorch_HiggsModel.py`. ### ### This tutorial parses the input model and runs the inference using ROOT's JITing capability. ### @@ -15,14 +15,15 @@ import ROOT # check if the input file exists -modelFile = "HiggsModel.keras" +modelFile = "HiggsModel.onnx" modelName = "HiggsModel" if not exists(modelFile): - raise FileNotFoundError("You need to run TMVA_Higgs_Classification.C to generate the Keras trained model") + raise FileNotFoundError("You need to run TMVA_SOFIE_PyTorch_HiggsModel.py to generate the ONNX trained model") -# parse the input Keras model into RModel object -model = ROOT.TMVA.Experimental.SOFIE.PyKeras.Parse(modelFile) +# parse the input ONNX model into RModel object +parser = ROOT.TMVA.Experimental.SOFIE.RModelParser_ONNX() +model = parser.Parse(modelFile) # generating inference code model.Generate() diff --git a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame_JIT.C b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame_JIT.C index 072384019ea51..97161ced30fb5 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_RDataFrame_JIT.C +++ b/tutorials/machine_learning/TMVA_SOFIE_RDataFrame_JIT.C @@ -1,10 +1,10 @@ /// \file /// \ingroup tutorial_ml /// \notebook -nodraw -/// This macro provides an example of using a trained model with Keras +/// This macro provides an example of using a trained model with PyTorch /// and make inference using SOFIE and RDataFrame -/// This macro uses as input a Keras model generated with the -/// TMVA_Higgs_Classification.C tutorial +/// This macro uses as input the SOFIE header generated from the ONNX model +/// with the TMVA_SOFIE_PyTorch_HiggsModel.py tutorial /// You need to run that macro before this one. /// In this case we are parsing the input file and then run the inference in the same /// macro making use of the ROOT JITing capability @@ -43,8 +43,9 @@ void TMVA_SOFIE_RDataFrame_JIT(std::string modelName = "HiggsModel"){ // check if the input file exists std::string modelHeaderFile = modelName + ".hxx"; if (gSystem->AccessPathName(modelHeaderFile.c_str())) { - Info("TMVA_SOFIE_RDataFrame","You need to run TMVA_SOFIE_Keras_Higgs_Model.py to generate the SOFIE header for the Keras trained model"); - return; + Info("TMVA_SOFIE_RDataFrame", "You need to run TMVA_SOFIE_PyTorch_HiggsModel.py to generate the SOFIE header " + "for the PyTorch trained model"); + return; } // check that also weigh file exists diff --git a/tutorials/machine_learning/TMVA_SOFIE_RSofieReader.C b/tutorials/machine_learning/TMVA_SOFIE_RSofieReader.C index 4da25f8afaea5..0a60c7662983b 100644 --- a/tutorials/machine_learning/TMVA_SOFIE_RSofieReader.C +++ b/tutorials/machine_learning/TMVA_SOFIE_RSofieReader.C @@ -1,16 +1,16 @@ /// \file /// \ingroup tutorial_ml /// \notebook -nodraw -/// This macro provides an example of using a trained model with Keras +/// This macro provides an example of using a trained model with PyTorch /// and make inference using SOFIE with the RSofieReader class -/// This macro uses as input a Keras model generated with the -/// TMVA_Higgs_Classification.C tutorial -/// You need to run that macro before to generate the trained Keras model +/// This macro uses as input an ONNX model generated with the +/// TMVA_SOFIE_PyTorch_HiggsModel.py tutorial +/// You need to run that macro before to generate the trained PyTorch model /// /// /// Execute in this order: /// ``` -/// root TMVA_Higgs_Classification.C +/// python3 TMVA_SOFIE_PyTorch_HiggsModel.py /// root TMVA_SOFIE_RSofieReader.C /// ``` /// @@ -22,9 +22,7 @@ using namespace TMVA::Experimental; void TMVA_SOFIE_RSofieReader(){ - RSofieReader model("HiggsModel.keras", {}, true ); - // for debugging - //RSofieReader model("Higgs_trained_model.keras", {}, true); + RSofieReader model("HiggsModel.onnx", {}, true); // the input shape for this model is a tensor with shape (1,7) diff --git a/tutorials/machine_learning/index.md b/tutorials/machine_learning/index.md index 242f934e21232..f6f4fd020bd98 100644 --- a/tutorials/machine_learning/index.md +++ b/tutorials/machine_learning/index.md @@ -118,15 +118,15 @@ | **Tutorial** || **Description** | |---------------|-----------------|-----------------| -| | TMVA_SOFIE_Inference.py | Using a trained model with Keras and make inference using SOFIE directly from Numpy. | +| | TMVA_SOFIE_Inference.py | Using a trained model with PyTorch and make inference using SOFIE directly from Numpy. | | TMVA_SOFIE_Keras.C | | Parsing of Keras .h5 file into RModel object and further generating the .hxx header files for inference. | -| TMVA_SOFIE_Keras_HiggsModel.C | | Run the SOFIE parser on the Keras model obtaining running TMVA_Higgs_Classification.C. You need to run that macro before this one. | -| | TMVA_SOFIE_Models.py | Inference with SOFIE using a set of models trained with Keras. | +| | TMVA_SOFIE_Models.py | Inference with SOFIE using a set of models trained with PyTorch. | | TMVA_SOFIE_ONNX.C | | Parsing of ONNX files into RModel object and further generating the .hxx header files for inference. | | | TMVA_SOFIE_PyTorch.py | Parsing of PyTorch .pt file into RModel object and further generating the .hxx header files for inference. | -| TMVA_SOFIE_RDataFrame.C | TMVA_SOFIE_RDataFrame.py | Inference with SOFIE and RDataFrame, of a model trained with Keras. | -| TMVA_SOFIE_RDataFrame_JIT.C | | Using a trained model with Keras and make inference using SOFIE and RDataFrame. | -| TMVA_SOFIE_RSofieReader.C | | Using a trained model with Keras and make inference using SOFIE with the RSofieReader class. | +| | TMVA_SOFIE_PyTorch_HiggsModel.py | Train a model on the Higgs dataset with PyTorch, export it to ONNX and run the SOFIE parser on it to generate the C++ inference code used by the other Higgs SOFIE tutorials. | +| TMVA_SOFIE_RDataFrame.C | TMVA_SOFIE_RDataFrame.py | Inference with SOFIE and RDataFrame, of a model trained with PyTorch. | +| TMVA_SOFIE_RDataFrame_JIT.C | | Using a trained model with PyTorch and make inference using SOFIE and RDataFrame. | +| TMVA_SOFIE_RSofieReader.C | | Using a trained model with PyTorch and make inference using SOFIE with the RSofieReader class. | \anchor data_loading ## Data loading for training