Add more flexibility to GCNFrameEncoder (#5)

* Enhance GCNFrameEncoder

* Added new GCN params to experiment and added random seed

Author: dobraczka

experiment.py

@@ -4,12 +4,15 @@
 import logging
 import os
 import pickle
+import random
 import shutil
 import time
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Type, get_args
 
 import click
+import numpy as np
+import torch
 from nephelai import upload
 from sylloge import OAEI, MovieGraphBenchmark, OpenEA
 from sylloge.base import EADataset
@@ -56,6 +59,15 @@
 logger = logging.getLogger("KlinkerExperiment")
 
 
+def set_random_seed(seed: Optional[int] = None):
+    if seed is None:
+        seed = np.random.randint(0, 2**16)
+        logger.info(f"No random seed provided. Using {seed}")
+    np.random.seed(seed=seed)
+    torch.manual_seed(seed=seed)
+    random.seed(seed)
+
+
 @dataclass
 class ExperimentInfo:
     params: Dict
@@ -194,21 +206,27 @@ def prepare(
 @click.option("--clean/--no-clean", default=True)
 @click.option("--wandb/--no-wandb", is_flag=True, default=False)
 @click.option("--nextcloud/--no-nextcloud", is_flag=True, default=False)
-def cli(clean: bool, wandb: bool, nextcloud: bool):
+@click.option("--random-seed", type=int, default=None)
+def cli(clean: bool, wandb: bool, nextcloud: bool, random_seed: Optional[int]):
     pass
 
 
 @cli.result_callback()
 def process_pipeline(
-    blocker_and_dataset: List, clean: bool, wandb: bool, nextcloud: bool
+    blocker_and_dataset: List,
+    clean: bool,
+    wandb: bool,
+    nextcloud: bool,
+    random_seed: Optional[int],
 ):
+    set_random_seed(random_seed)
     assert (
         len(blocker_and_dataset) == 2
     ), "Only 1 dataset and 1 blocker command can be used!"
     if not isinstance(blocker_and_dataset[0][0], EADataset):
         raise ValueError("First command must be dataset command!")
     if not isinstance(blocker_and_dataset[1][0], Blocker):
-        raise ValueError("First command must be blocker command!")
+        raise ValueError("Second command must be blocker command!")
     dataset_with_params, blocker_with_params = blocker_and_dataset
     dataset, ds_params = dataset_with_params
     blocker, bl_params, blocker_creation_time = blocker_with_params
@@ -609,20 +627,32 @@ def light_ea_blocker(
 
 @cli.command()
 @tokenized_frame_encoder_resolver.get_option(
-    "--inner-encoder", default="TransformerTokenizedFrameEncoder", as_string=True
+    "--inner-encoder", default="SIFEmbeddingTokenizedFrameEncoder", as_string=True
 )
+@click.option("--batch-size", type=int)
 @click.option("--embeddings", type=str, default="glove")
 @click.option("--depth", type=int, default=2)
-@click.option("--batch-size", type=int)
+@click.option("--edge-weight", type=float, default=1.0)
+@click.option("--self-loop-weight", type=float, default=2.0)
+@click.option("--layer-dims", type=int, default=300)
+@click.option("--bias", type=bool, default=True)
+@click.option("--use-weight-layers", type=bool, default=True)
+@click.option("--aggr", type=str, default="sum")
 @block_builder_resolver.get_option("--block-builder", default="kiez", as_string=True)
 @click.option("--block-builder-kwargs", type=str)
 @click.option("--n-neighbors", type=int, default=100)
 @click.option("--force", type=bool, default=True)
 def gcn_blocker(
     inner_encoder: Type[TokenizedFrameEncoder],
+    batch_size: Optional[int],
     embeddings: str,
     depth: int,
-    batch_size: Optional[int],
+    edge_weight: float,
+    self_loop_weight: float,
+    layer_dims: int,
+    bias: bool,
+    use_weight_layers: bool,
+    aggr: str,
     block_builder: Type[EmbeddingBlockBuilder],
     block_builder_kwargs: str,
     n_neighbors: int,
@@ -646,6 +676,12 @@ def gcn_blocker(
     blocker = EmbeddingBlocker(
         frame_encoder=GCNFrameEncoder(
             depth=depth,
+            edge_weight=edge_weight,
+            self_loop_weight=self_loop_weight,
+            layer_dims=layer_dims,
+            bias=bias,
+            use_weight_layers=use_weight_layers,
+            aggr=aggr,
             attribute_encoder=inner_encoder,
             attribute_encoder_kwargs=attribute_encoder_kwargs,
         ),

src/klinker/encoders/gcn.py

@@ -1,8 +1,10 @@
 import logging
-from typing import Optional, Tuple, Union
+import math
+from typing import List, Optional, Tuple, Union
 
 import numpy as np
 import torch
+import torch.nn as nn
 from class_resolver import HintOrType, OptionalKwargs
 
 try:
@@ -72,12 +74,11 @@ def _gcn_norm(
     edge_index,
     num_nodes: int,
     edge_weight=None,
-    improved=True,
+    fill_value=2.0,
     add_self_loops=True,
     flow="source_to_target",
     dtype=None,
 ):
-    fill_value = 2.0 if improved else 1.0
     assert flow in ["source_to_target", "target_to_source"]
 
     if edge_weight is None:
@@ -104,35 +105,119 @@ def _gcn_norm(
     return edge_index, edge_weight
 
 
+class BasicMessagePassing:
+    def __init__(
+        self,
+        edge_weight: float = 1.0,
+        self_loop_weight: float = 2.0,
+        aggr: str = "add",
+    ):
+        self.edge_weight = edge_weight
+        self.self_loop_weight = self_loop_weight
+        self.aggr = aggr
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
+        edge_index_with_loops, edge_weights = _gcn_norm(
+            edge_index,
+            num_nodes=len(x),
+            edge_weight=torch.tensor([self.edge_weight] * len(edge_index[0])),
+            fill_value=self.self_loop_weight,
+        )
+        return sparse_matmul(
+            SparseTensor.from_edge_index(edge_index_with_loops, edge_attr=edge_weights),
+            x,
+            reduce=self.aggr,
+        )
+
+
+def _glorot(value: torch.Tensor):
+    # see https://github.com/pyg-team/pytorch_geometric/blob/3e55a4c263f04ed6676618226f9a0aaf406d99b9/torch_geometric/nn/inits.py#L30
+    stdv = math.sqrt(6.0 / (value.size(-2) + value.size(-1)))
+    value.data.uniform_(-stdv, stdv)
+
+
+class FrozenGCNConv(BasicMessagePassing):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        bias: bool = False,
+        edge_weight: float = 1.0,
+        self_loop_weight: float = 2.0,
+        aggr: str = "add",
+    ):
+        super().__init__(
+            edge_weight=edge_weight, self_loop_weight=self_loop_weight, aggr=aggr
+        )
+        self.lin = nn.Linear(in_channels, out_channels, bias=bias)
+        for param in self.lin.parameters():
+            param.requires_grad = False
+        # Use glorot initialization
+        _glorot(self.lin.weight)
+
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
+        x = self.lin(x)
+        return super().forward(x, edge_index)
+
+
 class GCNFrameEncoder(RelationFrameEncoder):
     """Use untrained GCN for aggregating neighboring embeddings with self.
 
     Args:
         depth: How many hops of neighbors should be incorporated
+        edge_weight: Weighting of non-self-loops
+        self_loop_weight: Weighting of self-loops
+        layer_dims: Dimensionality of layers if used
+        bias: Whether to use bias in layers
+        use_weight_layers: Whether to use randomly initialized layers in aggregation
+        aggr: Which aggregation to use. Can be :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"` or :obj:`"max"`
         attribute_encoder: HintOrType[TokenizedFrameEncoder]: Base encoder class
         attribute_encoder_kwargs: OptionalKwargs: Keyword arguments for initializing encoder
     """
 
     def __init__(
         self,
         depth: int = 2,
+        edge_weight: float = 1.0,
+        self_loop_weight: float = 2.0,
+        layer_dims: int = 300,
+        bias: bool = False,
+        use_weight_layers: bool = True,
+        aggr: str = "sum",
         attribute_encoder: HintOrType[TokenizedFrameEncoder] = None,
         attribute_encoder_kwargs: OptionalKwargs = None,
     ):
         if not TORCH_SCATTER:
             logger.error("Could not find torch_scatter and/or torch_sparse package!")
         self.depth = depth
+        self.edge_weight = edge_weight
+        self.self_loop_weight = self_loop_weight
         self.device = resolve_device()
         self.attribute_encoder = tokenized_frame_encoder_resolver.make(
             attribute_encoder, attribute_encoder_kwargs
         )
-
-    def _forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
-        edge_index_with_loops, edge_weights = _gcn_norm(edge_index, num_nodes=len(x))
-        return sparse_matmul(
-            SparseTensor.from_edge_index(edge_index_with_loops, edge_attr=edge_weights),
-            x,
-        )
+        layers: List[BasicMessagePassing]
+        if use_weight_layers:
+            layers = [
+                FrozenGCNConv(
+                    in_channels=layer_dims,
+                    out_channels=layer_dims,
+                    edge_weight=edge_weight,
+                    self_loop_weight=self_loop_weight,
+                    aggr=aggr,
+                )
+                for _ in range(self.depth)
+            ]
+        else:
+            layers = [
+                BasicMessagePassing(
+                    edge_weight=edge_weight,
+                    self_loop_weight=self_loop_weight,
+                    aggr=aggr,
+                )
+                for _ in range(self.depth)
+            ]
+        self.layers = layers
 
     def _encode_rel(
         self,
@@ -143,6 +228,6 @@ def _encode_rel(
         full_graph = np.concatenate([rel_triples_left, rel_triples_right])
         edge_index = torch.from_numpy(full_graph[:, [0, 2]]).t()
         x = ent_features.vectors
-        for _ in range(self.depth):
-            x = self._forward(x, edge_index)
+        for layer in self.layers:
+            x = layer.forward(x, edge_index)
         return x

tests/test_blockers.py

@@ -21,14 +21,7 @@
     SimpleRelationalTokenBlocker,
     concat_neighbor_attributes,
 )
-from klinker.data import (
-    KlinkerBlockManager,
-    KlinkerDaskFrame,
-    KlinkerFrame,
-    KlinkerPandasFrame,
-    KlinkerTriplePandasFrame,
-    from_klinker_frame,
-)
+from klinker.data import KlinkerBlockManager, KlinkerFrame, from_klinker_frame
 from klinker.encoders.base import _get_ids
 
 
@@ -287,7 +280,12 @@ def test_assign_embedding_blocker(
 
 
 @pytest.mark.parametrize(
-    "cls, params", [("LightEAFrameEncoder", dict(mini_dim=3)), ("GCNFrameEncoder", {})]
+    "cls, params",
+    [
+        ("LightEAFrameEncoder", dict(mini_dim=3)),
+        ("GCNFrameEncoder", dict(layer_dims=3, use_weight_layers=True)),
+        ("GCNFrameEncoder", dict(layer_dims=3, use_weight_layers=False)),
+    ],
 )
 def test_assign_relation_frame_encoder(
     cls,