k3-node 1.0.0__py3-none-any.whl

k3_node/models/deep_graph_infomax.py

@@ -0,0 +1,113 @@
+import copy
+from typing import Callable, Optional
+
+import keras
+from keras import ops
+
+from k3_node.models.utils import reset, uniform_
+
+EPS = 1e-15
+
+
+class DeepGraphInfomax(keras.Model):
+    r"""The Deep Graph Infomax model from the
+    `"Deep Graph Infomax" <https://arxiv.org/abs/1809.10341>`_ paper.
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import GCNConv
+        from k3_node.models import DeepGraphInfomax
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        encoder = GCNConv(8, 16)
+        model = DeepGraphInfomax(
+            hidden_channels=16,
+            encoder=lambda x, edge_index: encoder(x, edge_index),
+            summary=lambda z, *args: keras.ops.sigmoid(keras.ops.mean(z, axis=0)),
+            corruption=lambda x, edge_index: (keras.random.shuffle(x), edge_index),  # permute node features
+        )
+        pos_z, neg_z, summary = model(x, edge_index)
+        print(tuple(pos_z.shape), tuple(neg_z.shape), tuple(summary.shape))  # (10, 16) (10, 16) (16,)
+        print(tuple(model.loss(pos_z, neg_z, summary).shape))  # (): scalar contrastive loss
+        ```
+    """
+
+    def __init__(
+        self,
+        hidden_channels: int,
+        encoder: Callable,
+        summary: Callable,
+        corruption: Optional[Callable] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_channels = hidden_channels
+        self.encoder = encoder
+        self.summary = summary
+        self.corruption = corruption if corruption is not None else default_corruption
+        self.weight = self.add_weight(
+            shape=(hidden_channels, hidden_channels),
+            initializer="glorot_uniform",
+            trainable=True,
+            name="weight",
+        )
+        self.built = True
+        self.loss = self.loss_fn
+
+    def reset_parameters(self):
+        reset(self.encoder)
+        reset(self.summary)
+        uniform_(self.hidden_channels, self.weight)
+
+    def call(self, *args, **kwargs):
+        r"""Returns the latent space for the input arguments, their
+        corruptions and their summary representation."""
+        if len(args) == 1 and isinstance(args[0], (tuple, list)):
+            args = tuple(args[0])
+        pos_z = self.encoder(*args, **kwargs)
+
+        cor = self.corruption(*args, **kwargs)
+        cor = cor if isinstance(cor, tuple) else (cor,)
+        cor_args = cor[: len(args)]
+        cor_kwargs = copy.copy(kwargs)
+        for key, value in zip(kwargs.keys(), cor[len(args):]):
+            cor_kwargs[key] = value
+
+        neg_z = self.encoder(*cor_args, **cor_kwargs)
+
+        summary = self.summary(pos_z, *args, **kwargs)
+
+        return pos_z, neg_z, summary
+
+    def discriminate(self, z, summary, sigmoid: bool = True):
+        r"""Given the patch-summary pair `z` and `summary`, computes the
+        probability scores assigned to this patch-summary pair."""
+        summary = ops.transpose(summary) if len(ops.shape(summary)) > 1 else summary
+        value = ops.matmul(z, ops.matmul(self.weight, summary))
+        return ops.sigmoid(value) if sigmoid else value
+
+    def loss_fn(self, pos_z, neg_z, summary):
+        r"""Computes the mutual information maximization objective."""
+        pos_loss = -ops.mean(ops.log(self.discriminate(pos_z, summary, sigmoid=True) + EPS))
+        neg_loss = -ops.mean(ops.log(1 - self.discriminate(neg_z, summary, sigmoid=True) + EPS))
+
+        return pos_loss + neg_loss
+
+    loss = loss_fn
+
+    def test(self, train_z, train_y, test_z, test_y, solver="lbfgs", *args, **kwargs):
+        r"""Evaluates latent space quality via a logistic regression
+        downstream task."""
+        from sklearn.linear_model import LogisticRegression
+
+        clf = LogisticRegression(*args, solver=solver, **kwargs).fit(
+            ops.convert_to_numpy(train_z), ops.convert_to_numpy(train_y)
+        )
+        return clf.score(ops.convert_to_numpy(test_z), ops.convert_to_numpy(test_y))
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.hidden_channels})"

k3_node/models/deepgcn.py

@@ -0,0 +1,121 @@
+import keras
+from keras import ops
+
+from k3_node.models.utils import reset
+
+
+class DeepGCNLayer(keras.layers.Layer):
+    r"""The skip connection operations from the
+    `"DeepGCNs: Can GCNs Go as Deep as CNNs?"
+    <https://arxiv.org/abs/1904.03751>`_ and `"All You Need to Train Deeper
+    GCNs" <https://arxiv.org/abs/2006.07739>`_ papers.
+    The implemented skip connections includes the pre-activation residual
+    connection (`"res+"`), the residual connection (`"res"`), the dense
+    connection (`"dense"`) and no connections (`"plain"`).
+
+    * **Res+** (`"res+"`):
+
+    .. math::
+        \text{Normalization}\to\text{Activation}\to\text{Dropout}\to
+        \text{GraphConv}\to\text{Res}
+
+    * **Res** (`"res"`) / **Dense** (`"dense"`) / **Plain** (`"plain"`):
+
+    .. math::
+        \text{GraphConv}\to\text{Normalization}\to\text{Activation}\to
+        \text{Res/Dense/Plain}\to\text{Dropout}
+
+    Args:
+        conv (optional): the GCN operator. (default: `None`)
+        norm (optional): the normalization layer. (default: `None`)
+        act (optional): the activation layer. (default: `None`)
+        block (str, optional): The skip connection operation to use
+            (`"res+"`, `"res"`, `"dense"` or `"plain"`). (default: `"res+"`)
+        dropout (float, optional): Whether to apply dropout. (default: `0.`)
+        ckpt_grad (bool, optional): Kept for API compatibility; a no-op,
+            since there is no gradient-checkpointing API shared uniformly
+            across Keras backends. (default: `False`)
+
+    Example:
+        ```python
+        import numpy as np
+        import keras
+        from k3_node.layers import GENConv, LayerNorm
+        from k3_node.models import DeepGCNLayer
+
+        x = np.random.rand(10, 8).astype("float32")  # 10 nodes with 8 features each
+        edge_index = np.random.randint(0, 10, size=(2, 30))  # 30 random edges
+
+        x = np.random.rand(10, 16).astype("float32")
+        layer = DeepGCNLayer(
+            conv=GENConv(16, 16),
+            norm=LayerNorm(16),
+            act=keras.layers.ReLU(),
+            block="res+",  # pre-activation residual block (DeeperGCN)
+        )
+        out = layer(x, edge_index)
+        print(tuple(out.shape))  # (10, 16)
+        ```
+    """
+    def __init__(
+        self,
+        conv=None,
+        norm=None,
+        act=None,
+        block: str = "res+",
+        dropout: float = 0.0,
+        ckpt_grad: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.conv = conv
+        self.norm = norm
+        self.act = act
+        self.block = block.lower()
+        assert self.block in ["res+", "res", "dense", "plain"]
+        self.dropout_rate = dropout
+        self.ckpt_grad = ckpt_grad
+        self.dropout = keras.layers.Dropout(dropout) if dropout > 0.0 else None
+        # No weights of its own; the wrapped layers build themselves on first use (PyG's examples
+        # sometimes only call `layer.conv` / `layer.norm` of a block).
+        self.built = True
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        reset(self.conv)
+        reset(self.norm)
+
+    def call(self, x, edge_index, *args, training=None, **kwargs):
+        if self.block == "res+":
+            h = x
+            if self.norm is not None:
+                h = self.norm(h, training=training)
+            if self.act is not None:
+                h = self.act(h)
+            if self.dropout is not None:
+                h = self.dropout(h, training=training)
+            h = self.conv(h, edge_index, *args, training=training, **kwargs)
+
+            return x + h
+
+        else:
+            h = self.conv(x, edge_index, *args, training=training, **kwargs)
+            if self.norm is not None:
+                h = self.norm(h, training=training)
+            if self.act is not None:
+                h = self.act(h)
+
+            if self.block == "res":
+                h = x + h
+            elif self.block == "dense":
+                h = ops.concatenate([x, h], axis=-1)
+            elif self.block == "plain":
+                pass
+
+            if self.dropout is not None:
+                h = self.dropout(h, training=training)
+            return h
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(block={self.block})"