Belle II Software light-2406-ragdoll
tensorflow.py
1#!/usr/bin/env python3
2
3
10
11import numpy as np
12import sys
13import os
14import tempfile
15
16
17class State:
18 """
19 Tensorflow state
20 """
21
22 def __init__(self, model=None, **kwargs):
23 """ Constructor of the state object """
24
25
26 self.model = model
27
28
29def feature_importance(state):
30 """
31 Return a list containing the feature importances
32 """
33 return []
34
35
36def get_model(number_of_features, number_of_spectators, number_of_events, training_fraction, parameters):
37 """
38 Return default tensorflow model
39 """
40 try:
41 import tensorflow as tf
42 except ImportError:
43 print("Please install tensorflow: pip3 install tensorflow")
44 sys.exit(1)
45
46 gpus = tf.config.list_physical_devices('GPU')
47 if gpus:
48 for gpu in gpus:
49 tf.config.experimental.set_memory_growth(gpu, True)
50
51 class my_model(tf.Module):
52
53 def __init__(self, **kwargs):
54
55 super().__init__(**kwargs)
56
57 self.W = tf.Variable(tf.ones(shape=(number_of_features, 1)), name="W")
58 self.b = tf.Variable(tf.ones(shape=(1, 1)), name="b")
59
60 self.optimizer = tf.optimizers.SGD(0.01)
61
62 @tf.function(input_signature=[tf.TensorSpec(shape=[None, number_of_features], dtype=tf.float32)])
63 def __call__(self, x):
64 return tf.nn.sigmoid(tf.matmul(self.clean_nans(x), self.W) + self.b)
65
66 def clean_nans(self, x):
67 return tf.where(tf.math.is_nan(x), tf.zeros_like(x), x)
68
69 def loss(self, predicted_y, target_y, w):
70 # cross entropy
71 epsilon = 1e-5
72 diff_from_truth = tf.where(target_y == 1., predicted_y, 1. - predicted_y)
73 return - tf.reduce_sum(w * tf.math.log(diff_from_truth + epsilon)) / tf.reduce_sum(w)
74
75 state = State(model=my_model())
76 return state
77
78
79def load(obj):
80 """
81 Load Tensorflow estimator into state
82 """
83 try:
84 import tensorflow as tf
85 except ImportError:
86 print("Please install tensorflow: pip3 install tensorflow")
87 sys.exit(1)
88
89 gpus = tf.config.list_physical_devices('GPU')
90 if gpus:
91 for gpu in gpus:
92 tf.config.experimental.set_memory_growth(gpu, True)
93
94 with tempfile.TemporaryDirectory() as path:
95
96 # recreate the expected folder structure
97 for subfolder in ['variables', 'assets']:
98 os.makedirs(os.path.join(path, subfolder))
99
100 file_names = obj[0]
101 for file_index, file_name in enumerate(file_names):
102 with open(f'{path}/{file_name}', 'w+b') as file:
103 file.write(bytes(obj[1][file_index]))
104
105 model = tf.saved_model.load(path)
106
107 state = State(model=model)
108 return state
109
110
111def apply(state, X):
112 """
113 Apply estimator to passed data.
114 """
115 try:
116 import tensorflow as tf
117 except ImportError:
118 print("Please install tensorflow: pip3 install tensorflow")
119 sys.exit(1)
120
121 r = state.model(tf.convert_to_tensor(np.atleast_2d(X), dtype=tf.float32)).numpy()
122 if r.shape[1] == 1:
123 r = r[:, 0] # cannot use squeeze because we might have output of shape [1,X classes]
124 return np.require(r, dtype=np.float32, requirements=['A', 'W', 'C', 'O'])
125
126
127def begin_fit(state, Xtest, Stest, ytest, wtest, nBatches):
128 """
129 Returns just the state object
130 """
131 state.nBatches = nBatches
132 return state
133
134
135def partial_fit(state, X, S, y, w, epoch, batch):
136 """
137 Pass batches of received data to tensorflow
138 """
139 try:
140 import tensorflow as tf
141 except ImportError:
142 print("Please install tensorflow: pip3 install tensorflow")
143 sys.exit(1)
144
145 with tf.GradientTape() as tape:
146 avg_cost = state.model.loss(state.model(X), y, w)
147 grads = tape.gradient(avg_cost, state.model.trainable_variables)
148
149 state.model.optimizer.apply_gradients(zip(grads, state.model.trainable_variables))
150
151 if batch == 0 and epoch == 0:
152 state.avg_costs = [avg_cost]
153 elif batch != state.nBatches-1:
154 state.avg_costs.append(avg_cost)
155 else:
156 # end of the epoch, print summary results, reset the avg_costs and update the counter
157 print(f"Epoch: {epoch:04d} cost= {np.mean(state.avg_costs):.9f}")
158 state.avg_costs = [avg_cost]
159
160 if epoch == 100000:
161 return False
162 return True
163
164
165def end_fit(state):
166 """
167 Store tensorflow model in a graph
168 """
169 try:
170 import tensorflow as tf
171 except ImportError:
172 print("Please install tensorflow: pip3 install tensorflow")
173 sys.exit(1)
174 with tempfile.TemporaryDirectory() as path:
175
176 tf.saved_model.save(state.model, path)
177 # tf.saved_model.save creates:
178 # path/saved_model.pb
179 # path/variables/variables.index
180 # path/variables/variables.data-00000-of-00001
181 # path/assets/* - This contains additional assets stored in the model.
182
183 file_names = ['saved_model.pb',
184 'variables/variables.index',
185 'variables/variables.data-00000-of-00001']
186
187 # we dont know what, if anything, is saved in assets/
188 assets_path = os.path.join(path, 'assets/')
189 file_names.extend([f'assets/{f.name}' for f in os.scandir(assets_path) if os.path.isfile(os.path.join(assets_path, f))])
190
191 files = []
192 for file_name in file_names:
193 with open(os.path.join(path, file_name), 'rb') as file:
194 files.append(file.read())
195 del state
196 return [file_names, files]
model
tensorflow model inheriting from tf.Module
Definition: tensorflow.py:26
def __init__(self, model=None, **kwargs)
Definition: tensorflow.py:22