feat(kbann): constraining version of kbann injector available

Author: MatteoMagnini

psyki/ski/__init__.py

@@ -45,9 +45,10 @@ def kins(model: Model,
     def kbann(model: Model,
               feature_mapping: dict[str, int],
               fuzzifier: str = 'towell',
-              omega: int = 4) -> Injector:
+              omega: float = 4.,
+              gamma: float = 10E-3) -> Injector:
         from psyki.ski.kbann import KBANN
-        return KBANN(model, feature_mapping, fuzzifier, omega)
+        return KBANN(model, feature_mapping, fuzzifier, omega, gamma)
 
 
 class EnrichedModel(Model):

psyki/ski/kbann/__init__.py

@@ -1,18 +1,28 @@
+import copy
 from typing import Iterable, Callable, List
+import tensorflow as tf
 from tensorflow.keras.layers import Concatenate
 from tensorflow import Tensor
+from tensorflow.keras.losses import Loss
+from tensorflow.python.keras.utils.generic_utils import custom_object_scope
 from psyki.logic.datalog.grammar import optimize_datalog_formula
 from psyki.logic import Fuzzifier, Formula
 from tensorflow.keras import Model
-from psyki.ski import Injector
+from psyki.ski import Injector, EnrichedModel
+from psyki.utils import model_deep_copy
 
 
 class KBANN(Injector):
     """
     Implementation of KBANN algorithm described by G. Towell in https://doi.org/10.1016/0004-3702(94)90105-8
     """
 
-    def __init__(self, predictor: Model, feature_mapping: dict[str, int], fuzzifier: str, omega: float = 4):
+    def __init__(self,
+                 predictor: Model,
+                 feature_mapping: dict[str, int],
+                 fuzzifier: str,
+                 omega: float = 4.,
+                 gamma: float = 10E-3):
         """
         @param predictor: the predictor.
         @param feature_mapping: a map between variables in the logic formulae and indices of dataset features. Example:
@@ -25,9 +35,46 @@ def __init__(self, predictor: Model, feature_mapping: dict[str, int], fuzzifier:
         """
         # self.feature_mapping: dict[str, int] = feature_mapping
         # Use as default fuzzifiers SubNetworkBuilder.
+        # TODO: analyse this warning that sometimes comes out, this should not be armful.
+        tf.get_logger().setLevel('ERROR')
         self._predictor = predictor
         self._fuzzifier = Fuzzifier.get(fuzzifier)([self._predictor.input, feature_mapping, omega])
         self._fuzzy_functions: Iterable[Callable] = ()
+        self.gamma = gamma
+
+    class ConstrainedModel(EnrichedModel):
+
+        def __init__(self, model: Model, gamma: float, custom_objects: dict):
+            super().__init__(model, custom_objects)
+            self.gamma = gamma
+            self.init_weights = copy.deepcopy(self.weights)
+
+        class CustomLoss(Loss):
+
+            def __init__(self, original_loss: Callable, model: Model, init_weights, gamma: float):
+                self.original_loss = original_loss
+                self.model = model
+                self.init_weights = init_weights
+                self.gamma = gamma
+                super().__init__()
+
+            def call(self, y_true, y_pred):
+                return self.original_loss(y_true, y_pred) + self.gamma * self._cost_factor()
+
+            def _cost_factor(self):
+                weights_quadratic_diff = 0
+                for init_weight, current_weight in zip(self.init_weights, self.model.weights):
+                    weights_quadratic_diff += tf.math.reduce_sum((init_weight - current_weight) ** 2)
+                # weights_quadratic_diff = tf.math.reduce_sum((tf.ragged.constant(self.init_weights) - tf.ragged.constant(self.weights)) ** 2)
+                return weights_quadratic_diff / (1 + weights_quadratic_diff)
+
+        def copy(self) -> EnrichedModel:
+            with custom_object_scope(self.custom_objects):
+                model = model_deep_copy(Model(self.input, self.output))
+                return KBANN.ConstrainedModel(model, self.gamma, self.custom_objects)
+
+        def loss_function(self, original_function: Callable) -> Callable:
+            return self.CustomLoss(original_function, self, self.init_weights, self.gamma)
 
     def inject(self, rules: List[Formula]) -> Model:
         # Prevent side effect on the original rules during optimization.
@@ -37,5 +84,5 @@ def inject(self, rules: List[Formula]) -> Model:
         predictor_input: Tensor = self._predictor.input
         modules = self._fuzzifier.visit(rules_copy)
         x = Concatenate(axis=1)(modules)
-        new_predictor = Model(predictor_input, x)
-        return self._fuzzifier.enriched_model(new_predictor)
+        #return self._fuzzifier.enriched_model(Model(predictor_input, x))
+        return self.ConstrainedModel(Model(predictor_input, x), self.gamma, self._fuzzifier.custom_objects)

test/psyki/injectors/test_injection.py

@@ -1,4 +1,5 @@
 import unittest
+from tensorflow.keras.losses import SparseCategoricalCrossentropy
 from psyki.ski import Injector
 from sklearn.datasets import load_iris
 from sklearn.model_selection import train_test_split, StratifiedKFold
@@ -93,7 +94,6 @@ class TestInjectionOnSpliceJunction(unittest.TestCase):
     x = get_binary_data(data.iloc[:, :-1], AGGREGATE_FEATURE_MAPPING)
     y.columns = [x.shape[1]]
     data = x.join(y)
-
     data, test = train_test_split(data, train_size=1000, random_state=0, stratify=data.iloc[:, -1])
     k_fold = StratifiedKFold(n_splits=10, shuffle=True, random_state=0)
     train_indices, _ = list(k_fold.split(data.iloc[:, :-1], data.iloc[:, -1:]))[0]
@@ -106,18 +106,28 @@ class TestInjectionOnSpliceJunction(unittest.TestCase):
     predictor = get_mlp(input_layer, 3, 3, [64, 32], 'relu', 'softmax', dropout=True)
     predictor = Model(input_layer, predictor)
 
-    def common_test_function(self, injector: Injector, batch_size: int, acceptable_accuracy: float):
+    def common_test_function(self, injector: Injector, batch_size: int, acceptable_accuracy: float, constrain=False):
         model = injector.inject(self.rules)
         # Test if clone is successful
         cloned_model = model.copy()
         del injector
 
-        model.compile('adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+        if constrain:
+            loss = model.loss_function(SparseCategoricalCrossentropy())
+        else:
+            loss = 'sparse_categorical_crossentropy'
+
+        model.compile('adam', loss=loss, metrics=['accuracy'])
         model.fit(self.train_x, self.train_y, batch_size=batch_size, epochs=self.EPOCHS, verbose=self.VERBOSE, callbacks=self.early_stop)
         accuracy = model.evaluate(self.test_x, self.test_y)[1]
         del model
 
-        cloned_model.compile('adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
+        if constrain:
+            loss = cloned_model.loss_function(SparseCategoricalCrossentropy())
+        else:
+            loss = 'sparse_categorical_crossentropy'
+
+        cloned_model.compile('adam', loss=loss, metrics=['accuracy'])
         cloned_model.fit(self.train_x, self.train_y, batch_size=batch_size, epochs=self.EPOCHS, verbose=self.VERBOSE, callbacks=self.early_stop)
         accuracy_cm = cloned_model.evaluate(self.test_x, self.test_y)[1]
         del cloned_model
@@ -127,7 +137,11 @@ def common_test_function(self, injector: Injector, batch_size: int, acceptable_a
 
     def test_kbann(self):
         injector = Injector.kbann(self.predictor, get_splice_junction_extended_feature_mapping(), 'towell', 1)
-        self.common_test_function(injector, batch_size=16, acceptable_accuracy=0.95)
+        self.common_test_function(injector, batch_size=16, acceptable_accuracy=0.957)
+
+    def test_kbann_with_constraining(self):
+        injector = Injector.kbann(self.predictor, get_splice_junction_extended_feature_mapping(), 'towell', 1, gamma=10E-5)
+        self.common_test_function(injector, batch_size=16, acceptable_accuracy=0.958, constrain=True)
 
     def test_kins(self):
         injector = Injector.kins(self.predictor, get_splice_junction_extended_feature_mapping())