Belle II Software light-2406-ragdoll
relations.py
1#!/usr/bin/env python3
2
3
10
11# This example serves as a basic example of implementing Relational networks into basf2 with tensorflow.
12# As a toy example it will try to tell if 2 out of multiple lines are hitting each other in three dimensional space.
13# Relevant Paper: https://arxiv.org/abs/1706.01427
14
15import tensorflow as tf
17import numpy as np
18
19
21 """ Using class to stop training early if it's not getting better"""
22
23 def __init__(self):
24 """ init class """
25
26 self.counter = 0
27
28 self.best_result = np.inf
29
30 def check(self, cost):
31 """
32 Check if validation result is better than the best validation result.
33 Decide if training should be continued.
34 """
35 if cost < self.best_result:
36 self.counter = 0
37 self.best_result = cost
38 else:
39 self.counter += 1
40 if self.counter >= 20:
41 return False
42 return True
43
44
45EARLY_STOPPER = early_stopping()
46
47
48def get_model(number_of_features, number_of_spectators, number_of_events, training_fraction, parameters):
49 """Building Graph inside tensorflow"""
50
51 gpus = tf.config.list_physical_devices('GPU')
52 if gpus:
53 for gpu in gpus:
54 tf.config.experimental.set_memory_growth(gpu, True)
55
56 class my_model(tf.Module):
57
58 def __init__(self, **kwargs):
59 super().__init__(**kwargs)
60
61 self.optimizer = tf.optimizers.Adam(0.01)
62 self.pre_optimizer = tf.optimizers.Adam(0.01)
63
64 def create_layer_variables(shape, name, activation_function):
65 weights = tf.Variable(
66 tf.random.truncated_normal(shape, stddev=1.0 / np.sqrt(float(shape[0]))),
67 name=f'{name}_weights')
68 biases = tf.Variable(tf.zeros(shape=[shape[1]]), name=f'{name}_biases')
69 return weights, biases, activation_function
70
71 n = int(number_of_features / 6)
72 self.number_of_track_pairs = int(n * (n - 1) / 2)
73 self.number_of_features_per_relation = 1
74 self.parameters = parameters
75
76 # build the relation net variables
77 self.rel_net_vars = []
78 shape = [12, 50, 50, 1]
79 for i in range(len(shape) - 1):
80 unit = tf.nn.tanh
81 if i == len(shape) - 2:
82 unit = tf.nn.sigmoid
83 self.rel_net_vars.append(create_layer_variables(shape[i:i + 2], f'relation_{i}', unit))
84
85 self.inference_vars = []
86
87 n_features_in = (self.number_of_track_pairs if self.parameters['use_relations'] else number_of_features)
88
89 if self.parameters['use_feed_forward']:
90 self.inference_vars.append(create_layer_variables([n_features_in, 50], 'inf_hidden_1', tf.nn.relu))
91 self.inference_vars.append(create_layer_variables([50, 50], 'inf_hidden_2', tf.nn.relu))
92 self.inference_vars.append(create_layer_variables([50, 1], 'inf_activation', tf.nn.sigmoid))
93 else:
94 self.inference_vars.append(create_layer_variables([n_features_in, 1], 'inf_activation', tf.nn.sigmoid))
95
96 def dense(self, x, W, b, activation_function):
97 return activation_function(tf.matmul(x, W) + b)
98
99 @tf.function(input_signature=[tf.TensorSpec(shape=[None, number_of_features], dtype=tf.float32)])
100 def prepare_x(self, x):
101 """ prepare the features for use in the relation net """
102 if not self.parameters['use_relations']:
103 return x
104
105 tracks = []
106 for i in range(int(number_of_features / 6)):
107 tracks.append(tf.slice(x, [0, i * 6], [-1, 6]))
108
109 relations = []
110
111 for feature_number in range(self.number_of_features_per_relation):
112 # Loop over every combination of input groups.
113 for counter, track1 in enumerate(tracks):
114 for track2 in tracks[counter + 1:]:
115 relations.append(self.relation_net(tf.concat([track1, track2], 1), self.rel_net_vars))
116
117 # workaround required for saving the model.
118 # Otherwise ValueErrors for trying to concatenate a list of length 0.
119 if relations:
120 return tf.concat(relations, 1)
121 else:
122 return x
123
124 @tf.function(input_signature=[tf.TensorSpec(shape=[None, 12], dtype=tf.float32)])
125 def pre_train(self, z):
126 """build net for pre-training with the same shared variables """
127 if not self.parameters['use_relations']:
128 return z
129
130 pre_training_relations = []
131 for feature_number in range(self.number_of_features_per_relation):
132 if self.parameters['pre_training_epochs'] > 0:
133 pre_training_relations.append(self.relation_net(z, self.rel_net_vars))
134
135 if pre_training_relations:
136 return tf.concat(pre_training_relations, 1)
137 else:
138 return z
139
140 def relation_net(self, x, variables):
141 """Build one relation net between 2 object using pre-build variables"""
142 for layer in variables:
143 x = self.dense(x, *layer)
144 return x
145
146 @tf.function(input_signature=[tf.TensorSpec(shape=[None, number_of_features], dtype=tf.float32)])
147 def __call__(self, x):
148 x = self.prepare_x(x)
149 for variables in self.inference_vars:
150 x = self.dense(x, *variables)
151 return x
152
153 @tf.function
154 def loss(self, predicted_y, target_y, w):
155 epsilon = 1e-5
156 diff_from_truth = tf.where(target_y == 1., predicted_y, 1. - predicted_y)
157 return - tf.reduce_sum(w * tf.math.log(diff_from_truth + epsilon)) / tf.reduce_sum(w)
158
159 state = State(model=my_model())
160 return state
161
162
163def begin_fit(state, Xtest, Stest, ytest, wtest, nBatches):
164 """Saves the training validation set for monitoring."""
165 state.val_x = Xtest
166 state.val_y = ytest
167 state.val_z = Stest
168
169 state.nBatches = nBatches
170 return state
171
172
173def partial_fit(state, X, S, y, w, epoch, batch):
174 """Pass received data to tensorflow session"""
175
176 def variable_is_trainable(var, parameters, pre_training=False):
177 out = True
178 if not parameters['use_relations'] and 'relation_' in var.name:
179 out = False
180 if pre_training and 'inf_' in var.name:
181 out = False
182 return out
183
184 # pre training trains shared variables on only 2 lines.
185 # In this case there is no relation net which have to compare two lines not hitting each other in a signal event.
186 if state.model.parameters['pre_training_epochs'] > epoch:
187 pre_trainable_variables = [
188 x for x in state.model.trainable_variables if variable_is_trainable(
189 x, state.model.parameters, True)]
190 with tf.GradientTape() as tape:
191 avg_cost = state.model.loss(state.model.pre_train(S), y, w)
192 grads = tape.gradient(avg_cost, pre_trainable_variables)
193 state.model.pre_optimizer.apply_gradients(zip(grads, pre_trainable_variables))
194
195 if state.epoch == epoch:
196 state.avg_costs.append()
197 else:
198 # started a new epoch, reset the avg_costs and update the counter
199 print(f"Pre-Training: Epoch: {epoch-1:05d}, cost={np.mean(state.avg_costs):.5f}")
200 state.avg_costs = [avg_cost]
201 state.epoch = epoch
202
203 # Training of the whole network.
204 else:
205 trainable_variables = [
206 x for x in state.model.trainable_variables if variable_is_trainable(
207 x, state.model.parameters, False)]
208 with tf.GradientTape() as tape:
209 avg_cost = state.model.loss(state.model(X), y, w)
210 grads = tape.gradient(avg_cost, trainable_variables)
211
212 state.model.optimizer.apply_gradients(zip(grads, trainable_variables))
213
214 if batch == 0 and epoch == 0:
215 state.avg_costs = [avg_cost]
216 elif batch != state.nBatches-1:
217 state.avg_costs.append(avg_cost)
218 else:
219 # started a new epoch, reset the avg_costs and update the counter
220 if epoch == state.model.parameters['pre_training_epochs']:
221 print(f"Pre-Training: Epoch: {epoch:05d}, cost={np.mean(state.avg_costs):.5f}")
222 else:
223 print(f"Epoch: {epoch:05d}, cost={np.mean(state.avg_costs):.5f}")
224
225 early_stopper_flag = EARLY_STOPPER.check(np.mean(state.avg_costs))
226 state.avg_costs = [avg_cost]
227 return early_stopper_flag
228 return True
229
230
231if __name__ == "__main__":
232 import os
233 import pandas
234 import uproot
235 import tempfile
236 import json
237
238 import basf2_mva
239 import basf2_mva_util
240 # ##############Building Data samples ###########################
241 # This is a dataset for testing relational nets.
242 # It consists of number_total_lines lines in 3 dimensional space.
243 # Each line has 6 variables.
244 # In apprx. half of the cases, two lines are hitting each other.
245 # This is considered a signal event.
246 # Training results differs from the number of total lines.
247
248 variables = []
249 # try using 10 lines and see what happens
250 number_total_lines = 5
251 # Names for the training data set
252 for i in range(number_total_lines):
253 variables += ['px_' + str(i), 'py_' + str(i), 'pz_' + str(i), 'dx_' + str(i), 'dy_' + str(i),
254 'dz_' + str(i)]
255 # Names for the spectator variables.
256 # Used as input variables for pre-training.
257 spectators = ['Spx1', 'Spy1', 'Spz1', 'Sdx1', 'Sdy1', 'Sdz1', 'Spx2', 'Spy2', 'Spz2', 'Sdx2', 'Sdy2', 'Sdz2']
258 # Number of events in training and test root file.
259 number_of_events = 100000
260
261 def build_signal_event():
262 """Building two lines which are hitting each other"""
263 p_vec1, p_vec2 = np.random.normal(size=3), np.random.normal(size=3)
264 v_cross = np.random.normal(size=3)
265 epsilon1, epsilon2 = np.random.rand() * 2 - 1, np.random.rand() * 2 - 1
266 v_vec1 = v_cross + (p_vec1 * epsilon1)
267 v_vec2 = v_cross + (p_vec2 * epsilon2)
268 return np.concatenate([p_vec1, v_vec1]), np.concatenate([p_vec2, v_vec2])
269
270 # This path will delete itself with all data in it after end of program.
271 with tempfile.TemporaryDirectory() as path:
272 for filename in ['train.root', 'test.root']:
273 print('Building ' + filename)
274 # Use random numbers to build all training and spectator variables.
275 data = np.random.normal(size=[number_of_events, number_total_lines * 6 + 12])
276 target = np.zeros([number_of_events], dtype=bool)
277
278 # Overwrite for half of the variables some lines so that they are hitting each other.
279 # Write them also at the end for the spectators.
280 for index, sample in enumerate(data):
281 if np.random.rand() > 0.5:
282 target[index] = True
283 i1, i2 = int(np.random.rand() * number_total_lines), int(np.random.rand() * (number_total_lines - 1))
284 i2 = (i1 + i2) % number_total_lines
285 track1, track2 = build_signal_event()
286 data[index, i1 * 6:(i1 + 1) * 6] = track1
287 data[index, i2 * 6:(i2 + 1) * 6] = track2
288 data[index, number_total_lines * 6:] = np.append(track1, track2)
289
290 # Saving all variables in root files
291 dic = {}
292 for i, name in enumerate(variables + spectators):
293 dic.update({name: data[:, i]})
294 dic.update({'isSignal': target})
295
296 df = pandas.DataFrame(dic, dtype=np.float32)
297 with uproot.recreate(os.path.join(path, filename)) as outfile:
298 outfile['variables'] = df
299
300 # ##########################Do Training#################################
301 # Do a comparison of different Nets for this task.
302
303 general_options = basf2_mva.GeneralOptions()
304 general_options.m_datafiles = basf2_mva.vector(os.path.join(path, 'train.root'))
305 general_options.m_treename = "variables"
306 general_options.m_variables = basf2_mva.vector(*variables)
307 general_options.m_target_variable = "isSignal"
308 general_options.m_spectators = basf2_mva.vector(*spectators)
309
310 specific_options = basf2_mva.PythonOptions()
311 specific_options.m_framework = "tensorflow"
312 specific_options.m_steering_file = 'mva/examples/tensorflow/relations.py'
313 specific_options.m_nIterations = 100
314 specific_options.m_mini_batch_size = 100
315
316 print('Train relational net with pre-training')
317 general_options.m_identifier = os.path.join(path, 'relation_2.xml')
318 specific_options.m_config = json.dumps({'use_relations': True, 'use_feed_forward': False, 'pre_training_epochs': 10})
319 basf2_mva.teacher(general_options, specific_options)
320
321 print('Train feed forward net')
322 general_options.m_identifier = os.path.join(path, 'feed_forward.xml')
323 specific_options.m_config = json.dumps({'use_relations': False, 'use_feed_forward': True, 'pre_training_epochs': 0})
324 basf2_mva.teacher(general_options, specific_options)
325
326 print('Train relational net')
327 general_options.m_identifier = os.path.join(path, 'relation.xml')
328 specific_options.m_config = json.dumps({'use_relations': True, 'use_feed_forward': True, 'pre_training_epochs': 0})
329 basf2_mva.teacher(general_options, specific_options)
330
331 # ########################Compare Results####################################
332
333 method1 = basf2_mva_util.Method(os.path.join(path, 'feed_forward.xml'))
334 method2 = basf2_mva_util.Method(os.path.join(path, 'relation.xml'))
335 method3 = basf2_mva_util.Method(os.path.join(path, 'relation_2.xml'))
336
337 test_data = basf2_mva.vector(os.path.join(path, 'test.root'))
338 print('Apply feed forward net')
339 p1, t1 = method1.apply_expert(test_data, general_options.m_treename)
340 print('Apply relational net')
341 p2, t2 = method2.apply_expert(test_data, general_options.m_treename)
342 print('Apply special relational net')
343 p3, t3 = method3.apply_expert(test_data, general_options.m_treename)
344
345 print('Feed Forward Net AUC: ', basf2_mva_util.calculate_auc_efficiency_vs_background_retention(p1, t1))
346 print('Relational Net AUC: ', basf2_mva_util.calculate_auc_efficiency_vs_background_retention(p2, t2))
347 print('Relational Net with pre-training AUC: ', basf2_mva_util.calculate_auc_efficiency_vs_background_retention(p3, t3))
def calculate_auc_efficiency_vs_background_retention(p, t, w=None)
best_result
saves best training result
Definition: relations.py:28
def check(self, cost)
Definition: relations.py:30
counter
counts how many times training is not getting better
Definition: relations.py:26