removed unused layers for fully model built to prevent errors in new pytorch trainer.

Author: PatReis

changelog.md

@@ -4,6 +4,7 @@ v4.0.1
 * Added further benchmark results for kgcnn version 4.
 * Fix error in ``kgcnn.layers.geom.PositionEncodingBasisLayer``
 * Fix error in ``kgcnn.literature.GCN.make_model_weighted``
+* Fix error in ``kgcnn.literature.AttentiveFP.make_model``
 * Had to change serialization for activation functions since with keras>=3.0.2 custom strings are not allowed also 
 causing clashes with built-in functions. We catch defaults to be at least backward compatible as possible and changed to serialization dictionary. Adapted all hyperparameter.
 * Renamed leaky_relu and swish in ``kgcnn.ops.activ`` to leaky_relu2 and swish2. 

kgcnn/layers/aggr.py

@@ -191,8 +191,9 @@ def __init__(self, pooling_method: str = "scatter_sum", pooling_index: int = glo
         self.normalize_by_weights = normalize_by_weights
         self.pooling_index = pooling_index
         self.pooling_method = pooling_method
-        self.to_aggregate = Aggregate(pooling_method=pooling_method)
-        self.to_aggregate_weights = Aggregate(pooling_method="scatter_sum")
+        # to_aggregate already made by super
+        if self.normalize_by_weights:
+            self.to_aggregate_weights = Aggregate(pooling_method="scatter_sum")
         self.axis_indices = axis_indices
 
     def build(self, input_shape):
@@ -201,7 +202,8 @@ def build(self, input_shape):
         node_shape, edges_shape, edge_index_shape, weights_shape = [list(x) for x in input_shape]
         edge_index_shape.pop(self.axis_indices)
         self.to_aggregate.build([tuple(x) for x in [edges_shape, edge_index_shape, node_shape]])
-        self.to_aggregate_weights.build([tuple(x) for x in [weights_shape, edge_index_shape, node_shape]])
+        if self.normalize_by_weights:
+            self.to_aggregate_weights.build([tuple(x) for x in [weights_shape, edge_index_shape, node_shape]])
         self.built = True
 
     def compute_output_shape(self, input_shape):

kgcnn/layers/attention.py

@@ -247,47 +247,97 @@ def get_config(self):
         return config
 
 
-class MultiHeadGATV2Layer(AttentionHeadGATV2):  # noqa
+class MultiHeadGATV2Layer(Layer):  # noqa
+    r"""Single layer for multiple Attention heads from :obj:`AttentionHeadGATV2` .
+
+    Uses concatenation or averaging of heads for final output.
+    """
 
     def __init__(self,
                  units: int,
                  num_heads: int,
                  activation: str = "kgcnn>leaky_relu2",
                  use_bias: bool = True,
                  concat_heads: bool = True,
+                 use_edge_features=False,
+                 use_final_activation=True,
+                 has_self_loops=True,
+                 kernel_regularizer=None,
+                 bias_regularizer=None,
+                 activity_regularizer=None,
+                 kernel_constraint=None,
+                 bias_constraint=None,
+                 kernel_initializer='glorot_uniform',
+                 bias_initializer='zeros',
+                 normalize_softmax: bool = False,
                  **kwargs):
-        super(MultiHeadGATV2Layer, self).__init__(
-            units=units,
-            activation=activation,
-            use_bias=use_bias,
-            **kwargs
-        )
+        r"""Initialize layer.
+
+        Args:
+            units (int): Units for the linear trafo of node features before attention.
+            num_heads: Number of attention heads.
+            concat_heads: Whether to concatenate heads or average.
+            use_edge_features (bool): Append edge features to attention computation. Default is False.
+            use_final_activation (bool): Whether to apply the final activation for the output.
+            has_self_loops (bool): If the graph has self-loops. Not used here. Default is True.
+            activation (str): Activation. Default is "kgcnn>leaky_relu2".
+            use_bias (bool): Use bias. Default is True.
+            kernel_regularizer: Kernel regularization. Default is None.
+            bias_regularizer: Bias regularization. Default is None.
+            activity_regularizer: Activity regularization. Default is None.
+            kernel_constraint: Kernel constrains. Default is None.
+            bias_constraint: Bias constrains. Default is None.
+            kernel_initializer: Initializer for kernels. Default is 'glorot_uniform'.
+            bias_initializer: Initializer for bias. Default is 'zeros'.
+        """
+        super(MultiHeadGATV2Layer, self).__init__(**kwargs)
         # Changes in keras serialization behaviour for activations in 3.0.2.
         # Keep string at least for default. Also renames to prevent clashes with keras leaky_relu.
         if activation in ["kgcnn>leaky_relu", "kgcnn>leaky_relu2"]:
             activation = {"class_name": "function", "config": "kgcnn>leaky_relu2"}
         self.num_heads = num_heads
         self.concat_heads = concat_heads
+        self.use_edge_features = use_edge_features
+        self.use_final_activation = use_final_activation
+        self.has_self_loops = has_self_loops
+        self.units = int(units)
+        self.normalize_softmax = normalize_softmax
+        self.use_bias = use_bias
+        kernel_args = {"kernel_regularizer": kernel_regularizer,
+                       "activity_regularizer": activity_regularizer, "bias_regularizer": bias_regularizer,
+                       "kernel_constraint": kernel_constraint, "bias_constraint": bias_constraint,
+                       "kernel_initializer": kernel_initializer, "bias_initializer": bias_initializer}
 
         self.head_layers = []
         for _ in range(num_heads):
-            lay_linear = Dense(units, activation=activation, use_bias=use_bias)
-            lay_alpha_activation = Dense(units, activation=activation, use_bias=use_bias)
-            lay_alpha = Dense(1, activation='linear', use_bias=False)
+            lay_linear = Dense(units, activation=activation, use_bias=use_bias, **kernel_args)
+            lay_alpha_activation = Dense(units, activation=activation, use_bias=use_bias, **kernel_args)
+            lay_alpha = Dense(1, activation='linear', use_bias=False, **kernel_args)
 
             self.head_layers.append((lay_linear, lay_alpha_activation, lay_alpha))
 
         self.lay_concat_alphas = Concatenate(axis=-2)
-        self.lay_concat_embeddings = Concatenate(axis=-2)
-        self.lay_pool_attention = AggregateLocalEdgesAttention()
-        # self.lay_pool = AggregateLocalEdges()
+
+        # self.lay_linear_trafo = Dense(units, activation="linear", use_bias=use_bias, **kernel_args)
+        # self.lay_alpha_activation = Dense(units, activation=activation, use_bias=use_bias, **kernel_args)
+        # self.lay_alpha = Dense(1, activation="linear", use_bias=False, **kernel_args)
+        self.lay_gather_in = GatherNodesIngoing()
+        self.lay_gather_out = GatherNodesOutgoing()
+        self.lay_concat = Concatenate(axis=-1)
+        self.lay_pool_attention = AggregateLocalEdgesAttention(normalize_softmax=normalize_softmax)
+        if self.use_final_activation:
+            self.lay_final_activ = Activation(activation=activation)
 
         if self.concat_heads:
             self.lay_combine_heads = Concatenate(axis=-1)
         else:
             self.lay_combine_heads = Average()
 
-    def __call__(self, inputs, **kwargs):
+    def build(self, input_shape):
+        """Build layer."""
+        super(MultiHeadGATV2Layer, self).build(input_shape)
+
+    def call(self, inputs, **kwargs):
         node, edge, edge_index = inputs
 
         # "a_ij" is a single-channel edge attention logits tensor. "a_ijs" is consequently the list which
@@ -338,6 +388,16 @@ def __call__(self, inputs, **kwargs):
     def get_config(self):
         """Update layer config."""
         config = super(MultiHeadGATV2Layer, self).get_config()
+        config.update({"use_edge_features": self.use_edge_features, "use_bias": self.use_bias,
+                       "units": self.units, "has_self_loops": self.has_self_loops,
+                       "normalize_softmax": self.normalize_softmax,
+                       "use_final_activation": self.use_final_activation})
+        if self.num_heads > 0:
+            conf_sub = self.head_layers[0][0].get_config()
+            for x in ["kernel_regularizer", "activity_regularizer", "bias_regularizer", "kernel_constraint",
+                      "bias_constraint", "kernel_initializer", "bias_initializer", "activation"]:
+                if x in conf_sub:
+                    config.update({x: conf_sub[x]})
         config.update({
             'num_heads': self.num_heads,
             'concat_heads': self.concat_heads

kgcnn/layers/conv.py

@@ -181,7 +181,6 @@ def call(self, inputs, **kwargs):
         Returns:
             Tensor: Updated node features.
         """
-        # print(inputs)
         node, edge, disjoint_indices = inputs
         x = self.lay_dense1(edge, **kwargs)
         x = self.lay_dense2(x, **kwargs)

kgcnn/literature/AttentiveFP/_make.py

@@ -119,7 +119,7 @@ def make_model(inputs: list = None,
         model_inputs,
         input_tensor_type=input_tensor_type,
         cast_disjoint_kwargs=cast_disjoint_kwargs,
-        mask_assignment=[0, 0, 1],
+        mask_assignment=[0, 1, 1],
         index_assignment=[None, None, 0]
     )
 

kgcnn/literature/DimeNetPP/_layers.py

@@ -129,7 +129,7 @@ def call(self, inputs, **kwargs):
         # Transform via 2D spherical basis
         sbf = self.dense_sbf1(sbf, **kwargs)
         sbf = self.dense_sbf2(sbf, **kwargs)
-        x_kj = self.lay_mult1([x_kj, sbf], **kwargs)
+        x_kj = self.lay_mult2([x_kj, sbf], **kwargs)
 
         # Aggregate interactions and up-project embeddings
         x_kj = self.lay_pool([rbf, x_kj, id_expand], **kwargs)

kgcnn/literature/HamNet/_layers.py

@@ -530,7 +530,7 @@ def call(self, inputs, **kwargs):
         q_u_ftr, q_v_ftr = self.gather_p([q_ftr, edi], **kwargs)
         p_u_ftr, p_v_ftr = self.gather_q([p_ftr, edi], **kwargs)
         p_uv_ftr = self.lazy_sub_p([p_v_ftr, p_u_ftr], **kwargs)
-        q_uv_ftr = self.lazy_sub_p([q_v_ftr, q_u_ftr], **kwargs)
+        q_uv_ftr = self.lazy_sub_q([q_v_ftr, q_u_ftr], **kwargs)
 
         attend_ftr = self.dense_attend(hv_v_ftr, **kwargs)
 

kgcnn/literature/MAT/_layers.py

@@ -40,6 +40,7 @@ class MATDistanceMatrix(ks.layers.Layer):
     def __init__(self, trafo: Union[str, None] = "exp", **kwargs):
         super(MATDistanceMatrix, self).__init__(**kwargs)
         self.trafo = trafo
+        # self._softmax = ks.layers.Softmax(axis=2)
         if self.trafo not in [None, "exp", "softmax"]:
             raise ValueError("`trafo` must be in [None, 'exp', 'softmax']")
 

kgcnn/literature/MEGAN/_make.py

@@ -1,3 +1,4 @@
+import keras as ks
 from ._model import MEGAN
 from kgcnn.models.utils import update_model_kwargs
 from kgcnn.layers.modules import Input

kgcnn/literature/MXMNet/_layers.py

@@ -119,7 +119,7 @@ def __init__(self, units: int = 64, output_units: int = 1, activation: str = "sw
 
         self.lin_rbf_out = Dense(self.dim, use_bias=False, activation="linear")
 
-        self.h_mlp = GraphMLP(self.dim, activation=activation)
+        # Fix for kgcnn==4.0.1: removed overwrite mlp here. Should not change model but prevents unused layers.
 
         self.y_mlp = GraphMLP([self.dim, self.dim, self.dim], activation=activation)
         self.y_W = Dense(self.output_dim, activation="linear",