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