Belle II Software light-2406-ragdoll
simple.py
1#!/usr/bin/env python3
2
3
10
11import time
12
14import torch
15
16
17class myModel(torch.nn.Module):
18 """
19 My dense neural network
20 """
21
22 def __init__(self, number_of_features):
23 """
24 Init the network
25 param: number_of_features number of input variables
26 """
27 super(myModel, self).__init__()
28
29
30 self.network = torch.nn.Sequential(
31 torch.nn.Linear(number_of_features, 128),
32 torch.nn.ReLU(),
33 torch.nn.Linear(128, 128),
34 torch.nn.ReLU(),
35 torch.nn.Linear(128, 1),
36 torch.nn.Sigmoid(),
37 )
38
39 def forward(self, x):
40 """
41 Run the network
42 """
43 prob = self.network(x)
44 return prob
45
46
47def get_model(number_of_features, number_of_spectators, number_of_events, training_fraction, parameters):
48 """
49 Returns default torch model
50 """
51
52 state = State(myModel(number_of_features).to("cpu"))
53 print(state.model)
54
55 state.optimizer = torch.optim.SGD(state.model.parameters(), parameters.get('learning_rate', 1e-2))
56
57 # we recreate the loss function on each batch so that we can pass in the weights
58 state.loss_fn = torch.nn.BCELoss
59
60 # for book keeping
61 state.epoch = 0
62 state.avg_costs = []
63 return state
64
65
66if __name__ == "__main__":
67 from basf2 import conditions
68 import basf2_mva
69 import basf2_mva_util
70 import json
71
72 # NOTE: do not use testing payloads in production! Any results obtained like this WILL NOT BE PUBLISHED
73 conditions.testing_payloads = [
74 'localdb/database.txt'
75 ]
76
77 general_options = basf2_mva.GeneralOptions()
78 general_options.m_datafiles = basf2_mva.vector("train.root")
79 general_options.m_identifier = "Simple"
80 general_options.m_treename = "tree"
81 variables = ['M', 'p', 'pt', 'pz',
82 'daughter(0, p)', 'daughter(0, pz)', 'daughter(0, pt)',
83 'daughter(1, p)', 'daughter(1, pz)', 'daughter(1, pt)',
84 'daughter(2, p)', 'daughter(2, pz)', 'daughter(2, pt)',
85 'chiProb', 'dr', 'dz',
86 'daughter(0, dr)', 'daughter(1, dr)',
87 'daughter(0, dz)', 'daughter(1, dz)',
88 'daughter(0, chiProb)', 'daughter(1, chiProb)', 'daughter(2, chiProb)',
89 'daughter(0, kaonID)', 'daughter(0, pionID)',
90 'daughterInvM(0, 1)', 'daughterInvM(0, 2)', 'daughterInvM(1, 2)']
91 general_options.m_variables = basf2_mva.vector(*variables)
92 general_options.m_target_variable = "isSignal"
93
94 specific_options = basf2_mva.PythonOptions()
95 specific_options.m_framework = "torch"
96 specific_options.m_steering_file = 'mva/examples/torch/simple.py'
97 # the number of training epochs
98 specific_options.m_nIterations = 64
99 specific_options.m_mini_batch_size = 256
100 specific_options.m_config = json.dumps({'learning_rate': 1e-2})
101 specific_options.m_training_fraction = 0.8
102 specific_options.m_normalise = False
103
104 training_start = time.time()
105 basf2_mva.teacher(general_options, specific_options)
106 training_stop = time.time()
107 training_time = training_stop - training_start
108 method = basf2_mva_util.Method(general_options.m_identifier)
109
110 inference_start = time.time()
111 test_data = ["test.root"]
112 p, t = method.apply_expert(basf2_mva.vector(*test_data), general_options.m_treename)
113 inference_stop = time.time()
114 inference_time = inference_stop - inference_start
116 print("Torch", training_time, inference_time, auc)
def calculate_auc_efficiency_vs_background_retention(p, t, w=None)
def forward(self, x)
Definition: simple.py:39
network
a dense model with one hidden layer
Definition: simple.py:30
def __init__(self, number_of_features)
Definition: simple.py:22