Logo ROOT  
Reference Guide
ApplicationRegressionPyTorch.py
Go to the documentation of this file.
1#!/usr/bin/env python
2## \file
3## \ingroup tutorial_tmva_pytorch
4## \notebook -nodraw
5## This tutorial shows how to apply a trained model to new data (regression).
6##
7## \macro_code
8##
9## \date 2020
10## \author Anirudh Dagar <anirudhdagar6@gmail.com> - IIT, Roorkee
11
12
13from ROOT import TMVA, TFile, TString
14from array import array
15from subprocess import call
16from os.path import isfile
17
18
19# Setup TMVA
22reader = TMVA.Reader("Color:!Silent")
23
24
25# Load data
26if not isfile('tmva_reg_example.root'):
27 call(['curl', '-L', '-O', 'http://root.cern.ch/files/tmva_reg_example.root'])
28
29data = TFile.Open('tmva_reg_example.root')
30tree = data.Get('TreeR')
31
32branches = {}
33for branch in tree.GetListOfBranches():
34 branchName = branch.GetName()
35 branches[branchName] = array('f', [-999])
36 tree.SetBranchAddress(branchName, branches[branchName])
37 if branchName != 'fvalue':
38 reader.AddVariable(branchName, branches[branchName])
39
40
41# Book methods
42reader.BookMVA('PyTorch', TString('dataset/weights/TMVARegression_PyTorch.weights.xml'))
43
44
45# Define predict function
46def predict(model, test_X, batch_size=32):
47 # Set to eval mode
48 model.eval()
49
50 test_dataset = torch.utils.data.TensorDataset(torch.Tensor(test_X))
51 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
52
53 predictions = []
54 with torch.no_grad():
55 for i, data in enumerate(test_loader):
56 X = data[0]
57 outputs = model(X)
58 predictions.append(outputs)
59 preds = torch.cat(predictions)
60
61 return preds.numpy()
62
63load_model_custom_objects = {"optimizer": None, "criterion": None, "train_func": None, "predict_func": predict}
64
65
66# Print some example regressions
67print('Some example regressions:')
68for i in range(20):
69 tree.GetEntry(i)
70 print('True/MVA value: {}/{}'.format(branches['fvalue'][0],reader.EvaluateMVA('PyTorch')))
static TFile * Open(const char *name, Option_t *option="", const char *ftitle="", Int_t compress=ROOT::RCompressionSetting::EDefaults::kUseCompiledDefault, Int_t netopt=0)
Create / open a file.
Definition: TFile.cxx:4013
static void PyInitialize()
Initialize Python interpreter.
The Reader class serves to use the MVAs in a specific analysis context.
Definition: Reader.h:64
static Tools & Instance()
Definition: Tools.cxx:71
Basic string class.
Definition: TString.h:136
def predict(model, test_X, batch_size=100)