Merge pull request #605 from NNPDF/tf2.0.0-hotfix

TF 2 compatibility

conda-recipe/meta.yaml

@@ -23,7 +23,7 @@ requirements:
         - numpy
     run:
         - tensorflow
-        - keras
+        - keras >=2.3 # 2.2.4 is not compatible with tf 2.0 while tf 2.0 only forces 2.2.4 ...
         - hyperopt
         - seaborn
         - lhapdf

n3fit/src/n3fit/ModelTrainer.py

@@ -409,11 +409,7 @@ def _train_and_fit(self, stopping_object, epochs):
 
             if (epoch + 1) % 100 == 0:
                 print_stats = True
-                for pos_ds in self.training["posdatasets"]:
-                    name = pos_ds
-                    curr_w = training_model.get_layer(name).get_weights()
-                    new_w = [curr_w[0] * pos_multiplier]
-                    training_model.get_layer(name).set_weights(new_w)
+                training_model.multiply_weights('pos', self.training["posdatasets"], pos_multiplier)
 
             passes = stopping_object.monitor_chi2(out, epoch, print_stats=print_stats)
 

n3fit/src/n3fit/backends/__init__.py

@@ -1,3 +1,5 @@
+import tensorflow as tf
+tf.compat.v1.disable_v2_behavior()
 from n3fit.backends.keras_backend.internal_state import set_initial_state, clear_backend_state
 from n3fit.backends.keras_backend.MetaLayer import MetaLayer
 from n3fit.backends.keras_backend.MetaModel import MetaModel

n3fit/src/n3fit/backends/keras_backend/MetaLayer.py

@@ -8,7 +8,7 @@
     For instance: np_to_tensor is just a call to K.constant
 """
 
-
+import tensorflow as tf
 from keras import backend as K
 from keras.engine.topology import Layer
 from keras.initializers import Constant, RandomUniform, glorot_normal, glorot_uniform
@@ -113,7 +113,7 @@ def tensor_product(self, tensor_x, tensor_y, axes, **kwargs):
         """
         Computes the tensordot product between tensor_x and tensor_y
         """
-        return K.tf.tensordot(tensor_x, tensor_y, axes=axes, **kwargs)
+        return tf.tensordot(tensor_x, tensor_y, axes=axes, **kwargs)
 
     def transpose(self, tensor, **kwargs):
         """
@@ -125,7 +125,7 @@ def boolean_mask(self, tensor, mask, axis=None, **kwargs):
         """
         Applies boolean mask to a tensor
         """
-        return K.tf.boolean_mask(tensor, mask, axis=axis, **kwargs)
+        return tf.boolean_mask(tensor, mask, axis=axis, **kwargs)
 
     def concatenate(self, tensor_list, axis=-1, target_shape=None):
         """

n3fit/src/n3fit/backends/keras_backend/MetaModel.py

@@ -4,7 +4,7 @@
     Extension of the backend Model class containing some wrappers in order to absorb other
     backend-dependent calls
 """
-from keras.models import Model
+from keras.models import Model, Sequential
 import keras.optimizers as Kopt
 
 from n3fit.backends.keras_backend.operations import numpy_to_input
@@ -67,6 +67,10 @@ def __init__(self, input_tensors, output_tensors, extra_tensors=None, **kwargs):
                 input_list += inputs
                 output_list.append(o_tensor)
 
+        self.all_inputs = input_list
+        self.all_outputs = output_list
+        self.internal_models = {}
+
         super(MetaModel, self).__init__(input_list, output_list, **kwargs)
 
     def fit(self, epochs=1, **kwargs):
@@ -156,3 +160,23 @@ def compile(
         super(MetaModel, self).compile(
             optimizer=opt, loss=loss, target_tensors=target_output, **kwargs
         )
+
+    def multiply_weights(self, key, layer_names, multiplier):
+        """ Multiply all weights for the given layers by some scalar
+
+        Parameters
+        ----------
+            layer_names: list of names of the layers to update weights
+            multiplier: scalar number to multiply the weights with
+        """
+        internal_model = self.internal_models.get(key)
+        if not internal_model:
+            # Create an internal model to access the weights of the
+            # layer we want to update
+            # is this a bug or a feature?
+            layers = [self.get_layer(i) for i in layer_names]
+            internal_model = Sequential(layers)
+            self.internal_models[key] = internal_model
+        current_weights = internal_model.get_weights()
+        new_weights = [i*multiplier for i in current_weights]
+        internal_model.set_weights(new_weights)

n3fit/src/n3fit/backends/keras_backend/internal_state.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import random as rn
+import tensorflow as tf
 from keras import backend as K
 
 
@@ -25,11 +26,11 @@ def set_initial_state(seed=13):
     # Clear the state of keras in case anyone used it before
     K.clear_session()
 
-    session_conf = K.tf.ConfigProto(intra_op_parallelism_threads=1,
+    session_conf = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=1,
                                 inter_op_parallelism_threads=1)
-    K.tf.set_random_seed(use_seed)
-    sess = K.tf.Session(graph=K.tf.get_default_graph(), config=session_conf)
-    K.set_session(sess)
+    tf.compat.v1.set_random_seed(use_seed)
+    sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=session_conf)
+    tf.compat.v1.keras.backend.set_session(sess)
 
     return 0
 
@@ -45,7 +46,7 @@ def clear_backend_state():
 
     K.clear_session()
     # Don't open threads that you are not going to eat
-    session_conf = K.tf.ConfigProto(intra_op_parallelism_threads=2,
+    session_conf = tf.compat.v1.ConfigProto(intra_op_parallelism_threads=2,
                             inter_op_parallelism_threads=8)
-    sess = K.tf.Session(graph=K.tf.get_default_graph(), config=session_conf)
-    K.set_session(sess)
+    sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=session_conf)
+    tf.compat.v1.keras.backend.set_session(sess)

n3fit/src/n3fit/backends/keras_backend/losses.py

@@ -2,6 +2,7 @@
     Module containing a list of loss functions availables to the fitting code
 """
 
+import tensorflow as tf
 import keras.backend as K
 
 
@@ -15,8 +16,8 @@ def l_invcovmat(invcovmat_np):
     def true_loss(y_true, y_pred):
         # (yt - yp) * covmat * (yt - yp)
         tmp = y_true - y_pred
-        right_dot = K.tf.tensordot(invcovmat, K.transpose(tmp), axes=1)
-        return K.tf.tensordot(tmp, right_dot, axes=1)
+        right_dot = tf.tensordot(invcovmat, K.transpose(tmp), axes=1)
+        return tf.tensordot(tmp, right_dot, axes=1)
 
     return true_loss