Logo ROOT  
Reference Guide
 
Loading...
Searching...
No Matches
ApplicationClassificationPyTorch.py
Go to the documentation of this file.
1## \file
2## \ingroup tutorial_tmva_pytorch
3## \notebook -nodraw
4## This tutorial shows how to apply a trained model to new data.
5##
6## \macro_code
7##
8## \date 2020
9## \author Anirudh Dagar <anirudhdagar6@gmail.com> - IIT, Roorkee
10
11
12# PyTorch has to be imported before ROOT to avoid crashes because of clashing
13# std::regexp symbols that are exported by cppyy.
14# See also: https://github.com/wlav/cppyy/issues/227
15import torch
16
17from ROOT import TMVA, TFile, TString, gROOT
18from array import array
19from subprocess import call
20from os.path import isfile
21
22
23# Setup TMVA
26reader = TMVA.Reader("Color:!Silent")
27
28
29# Load data
30fname = str(gROOT.GetTutorialDir()) + "/machine_learning/data/tmva_class_example.root"
31data = TFile.Open(fname)
32signal = data.Get('TreeS')
33background = data.Get('TreeB')
34
35branches = {}
36for branch in signal.GetListOfBranches():
37 branchName = branch.GetName()
38 branches[branchName] = array('f', [-999])
39 reader.AddVariable(branchName, branches[branchName])
40 signal.SetBranchAddress(branchName, branches[branchName])
41 background.SetBranchAddress(branchName, branches[branchName])
42
43
44# Define predict function
45def predict(model, test_X, batch_size=32):
46 # Set to eval mode
48
49 test_dataset = torch.utils.data.TensorDataset(torch.Tensor(test_X))
50 test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
51
52 predictions = []
53 with torch.no_grad():
54 for i, data in enumerate(test_loader):
55 X = data[0]
56 outputs = model(X)
57 predictions.append(outputs)
58 preds = torch.cat(predictions)
59
60 return preds.numpy()
61
62
63load_model_custom_objects = {"optimizer": None, "criterion": None, "train_func": None, "predict_func": predict}
64
65
66# Book methods
67reader.BookMVA('PyTorch', TString('dataset/weights/TMVAClassification_PyTorch.weights.xml'))
68
69
70# Print some example classifications
71print('Some signal example classifications:')
72for i in range(20):
74 print(reader.EvaluateMVA('PyTorch'))
75print('')
76
77print('Some background example classifications:')
78for i in range(20):
80 print(reader.EvaluateMVA('PyTorch'))
ROOT::Detail::TRangeCast< T, true > TRangeDynCast
TRangeDynCast is an adapter class that allows the typed iteration through a TCollection.
The Reader class serves to use the MVAs in a specific analysis context.
Definition Reader.h:64
Basic string class.
Definition TString.h:139