Logo ROOT  
Reference Guide
ApplicationClassificationPyTorch.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.
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_class_example.root'):
27 call(['curl', '-L', '-O', 'http://root.cern.ch/files/tmva_class_example.root'])
28
29data = TFile.Open('tmva_class_example.root')
30signal = data.Get('TreeS')
31background = data.Get('TreeB')
32
33branches = {}
34for branch in signal.GetListOfBranches():
35 branchName = branch.GetName()
36 branches[branchName] = array('f', [-999])
37 reader.AddVariable(branchName, branches[branchName])
38 signal.SetBranchAddress(branchName, branches[branchName])
39 background.SetBranchAddress(branchName, branches[branchName])
40
41
42# Define predict function
43def predict(model, test_X, batch_size=32):
44 # Set to eval mode
45 model.eval()
46
47 test_dataset = torch.utils.data.TensorDataset(torch.Tensor(test_X))
48 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
49
50 predictions = []
51 with torch.no_grad():
52 for i, data in enumerate(test_loader):
53 X = data[0]
54 outputs = model(X)
55 predictions.append(outputs)
56 preds = torch.cat(predictions)
57
58 return preds.numpy()
59
60
61load_model_custom_objects = {"optimizer": None, "criterion": None, "train_func": None, "predict_func": predict}
62
63
64# Book methods
65reader.BookMVA('PyTorch', TString('dataset/weights/TMVAClassification_PyTorch.weights.xml'))
66
67
68# Print some example classifications
69print('Some signal example classifications:')
70for i in range(20):
71 signal.GetEntry(i)
72 print(reader.EvaluateMVA('PyTorch'))
73print('')
74
75print('Some background example classifications:')
76for i in range(20):
77 background.GetEntry(i)
78 print(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)