Upload 3 files

models/acm_gin.py

@@ -0,0 +1,207 @@
+"""
+(c) Adaptation of the code from https://github.com/SitaoLuan/ACM-GNN
+"""
+
+import torch
+import torch.nn as nn
+
+from torch import Tensor
+from typing import Union
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.inits import reset
+from torch_geometric.typing import OptPairTensor, OptTensor, Size
+from torch_geometric.utils import scatter
+
+from .utils import create_activation
+
+
+class ACM_GIN(MessagePassing):
+    def __init__(
+        self,
+        nn_lowpass: torch.nn.Module,
+        nn_highpass: torch.nn.Module,
+        nn_fullpass: torch.nn.Module,
+        nn_lowpass_proj: torch.nn.Module,
+        nn_highpass_proj: torch.nn.Module,
+        nn_fullpass_proj: torch.nn.Module,
+        nn_mix: torch.nn.Module,
+        T: float = 3.0,
+        **kwargs,
+    ):
+        kwargs.setdefault("aggr", "add")
+        super().__init__(**kwargs)
+        self.nn_lowpass = nn_lowpass
+        self.nn_highpass = nn_highpass
+        self.nn_fullpass = nn_fullpass
+        self.nn_lowpass_proj = nn_lowpass_proj
+        self.nn_highpass_proj = nn_highpass_proj
+        self.nn_fullpass_proj = nn_fullpass_proj
+        self.nn_mix = nn_mix
+        self.sigmoid = torch.nn.Sigmoid()
+        self.softmax = torch.nn.Softmax(dim=1)
+        self.T = T
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        reset(self.nn_lowpass)
+        reset(self.nn_highpass)
+        reset(self.nn_fullpass)
+        reset(self.nn_lowpass_proj)
+        reset(self.nn_highpass_proj)
+        reset(self.nn_fullpass_proj)
+        reset(self.nn_mix)
+
+    def forward(
+        self,
+        x: Union[Tensor, OptPairTensor],
+        edge_index: Tensor,
+        edge_weight: OptTensor = None,
+        size: Size = None,
+    ) -> Tensor:
+
+        if isinstance(x, Tensor):
+            x: OptPairTensor = (x, x)
+
+        # propagate_type: (x: OptPairTensor, edge_attr: OptTensor)
+        out = self.propagate(edge_index, x=x, edge_weight=edge_weight, size=size)
+
+        deg = scatter(edge_weight, edge_index[1], 0, out.size(0), reduce="sum")
+        deg_inv = 1.0 / deg
+        deg_inv.masked_fill_(deg_inv == float("inf"), 0)
+        out = deg_inv.view(-1, 1) * out
+
+        x_r = x[1]
+        if x_r is not None:
+            out_lowpass = (x_r + out) / 2.0
+            out_highpass = (x_r - out) / 2.0
+
+        # compute embeddings for each filter
+        out_lowpass = self.nn_lowpass(out_lowpass)
+        out_highpass = self.nn_highpass(out_highpass)
+        out_fullpass = self.nn_fullpass(x_r)
+        # compute importance weights per filter
+        alpha_lowpass = self.sigmoid(self.nn_lowpass_proj(out_lowpass))
+        alpha_highpass = self.sigmoid(self.nn_highpass_proj(out_highpass))
+        alpha_fullpass = self.sigmoid(self.nn_fullpass_proj(out_fullpass))
+        alpha_cat = torch.concat([alpha_lowpass, alpha_highpass, alpha_fullpass], dim=1)
+        alpha_cat = self.softmax(self.nn_mix(alpha_cat / self.T))
+
+        out = alpha_cat[:, 0].view(-1, 1) * out_lowpass
+        out = out + alpha_cat[:, 1].view(-1, 1) * out_highpass
+        out = out + alpha_cat[:, 2].view(-1, 1) * out_fullpass
+
+        return out
+
+    def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
+        return edge_weight.view(-1, 1) * x_j
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(nn={self.nn})"
+
+
+class ACM_GIN_model(nn.Module):
+    """ """
+
+    def __init__(
+        self, in_dim, out_dim, num_layers, hidden_dim, batchnorm, activation="relu"
+    ):
+        super(ACM_GIN_model, self).__init__()
+        self.num_layers = num_layers
+        self.hidden_dim = hidden_dim
+        self.gnn_batchnorm = batchnorm
+        self.out_dim = out_dim
+
+        self.ACM_convs = nn.ModuleList()
+        self.nns_lowpass = nn.ModuleList()
+        self.nns_highpass = nn.ModuleList()
+        self.nns_fullpass = nn.ModuleList()
+        self.nns_lowpass_proj = nn.ModuleList()
+        self.nns_highpass_proj = nn.ModuleList()
+        self.nns_fullpass_proj = nn.ModuleList()
+        self.nns_mix = nn.ModuleList()
+
+        self.activation = create_activation(activation)
+
+        for i in range(self.num_layers):
+            # projection modules to compute importance weights
+            for channel_proj_module in [
+                self.nns_lowpass_proj,
+                self.nns_highpass_proj,
+                self.nns_fullpass_proj,
+            ]:
+                if i == self.num_layers - 1:
+                    channel_proj_module.append(nn.Linear(self.out_dim, 1))
+                else:
+                    channel_proj_module.append(nn.Linear(self.hidden_dim, 1))
+            # weights mixing module as attention mechanism
+            self.nns_mix.append(nn.Linear(3, 3))
+
+            # GIN embedding scheme per channel
+            if i == 0:
+                local_input_dim = in_dim
+            else:
+                local_input_dim = self.hidden_dim
+
+            if i == self.num_layers - 1:
+                local_out_dim = self.out_dim
+            else:
+                local_out_dim = self.hidden_dim
+
+            for channel_module in [
+                self.nns_lowpass,
+                self.nns_highpass,
+                self.nns_fullpass,
+            ]:
+                if self.gnn_batchnorm:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        nn.BatchNorm1d(self.hidden_dim),
+                        self.activation,
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        nn.BatchNorm1d(local_out_dim),
+                        self.activation,
+                    )
+                else:
+                    sequential = nn.Sequential(
+                        nn.Linear(local_input_dim, self.hidden_dim),
+                        self.activation,
+                        nn.Linear(self.hidden_dim, local_out_dim),
+                        self.activation,
+                    )
+
+                channel_module.append(sequential)
+
+            self.ACM_convs.append(
+                ACM_GIN(
+                    nn_lowpass=self.nns_lowpass[i],
+                    nn_highpass=self.nns_highpass[i],
+                    nn_fullpass=self.nns_fullpass[i],
+                    nn_lowpass_proj=self.nns_lowpass_proj[i],
+                    nn_highpass_proj=self.nns_highpass_proj[i],
+                    nn_fullpass_proj=self.nns_fullpass_proj[i],
+                    nn_mix=self.nns_mix[i],
+                )
+            )
+
+    def reset_parameters(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                m.reset_parameters()
+            elif isinstance(m, nn.BatchNorm1d):
+                m.reset_parameters()
+
+    def forward(self, x, edge_index, edge_attr, return_hidden=False):
+        outs = []
+        for i in range(self.num_layers):
+            x = self.ACM_convs[i](x=x, edge_index=edge_index, edge_weight=edge_attr)
+            outs.append(x)
+        if return_hidden:
+            return x, outs
+        else:
+            return x
+
+
+if __name__ == "__main__":
+    acm_gin = ACM_GIN_model(46, 46, 2, 256, True)
+    print(sum(p.numel() for p in acm_gin.parameters() if p.requires_grad))
+    print("")

models/edcoder.py

@@ -0,0 +1,261 @@
+"""
+(c) Adaptation of the code from https://github.com/THUDM/GraphMAE
+"""
+
+from typing import Optional
+from itertools import chain
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.utils import dropout_edge
+from torch_geometric.utils import add_self_loops
+
+from .acm_gin import ACM_GIN_model
+
+
+def sce_loss(x, y, alpha=3):
+    x = F.normalize(x, p=2, dim=-1)
+    y = F.normalize(y, p=2, dim=-1)
+
+    loss = (1 - (x * y).sum(dim=-1)).pow_(alpha)
+    loss = loss.mean()
+
+    return loss
+
+
+def setup_module(
+    m_type,
+    in_dim,
+    out_dim,
+    num_hidden,
+    num_layers,
+    activation,
+    batchnorm,
+) -> nn.Module:
+
+    if m_type == "acm_gin":
+        mod = ACM_GIN_model(
+            int(in_dim),
+            int(out_dim),
+            num_layers,
+            int(num_hidden),
+            batchnorm,
+            activation=activation,
+        )
+    else:
+        raise NotImplementedError
+
+    return mod
+
+
+class PreModel(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        edge_in_dim: int,
+        num_hidden: int,
+        num_layers: int,
+        nhead: int,
+        nhead_out: int,
+        activation: str,
+        feat_drop: float,
+        attn_drop: float,
+        negative_slope: float,
+        residual: bool,
+        norm: Optional[str],
+        mask_rate: float = 0.3,
+        encoder_type: str = "gat",
+        decoder_type: str = "gat",
+        loss_fn: str = "sce",
+        drop_edge_rate: float = 0.0,
+        replace_rate: float = 0.1,
+        alpha_l: float = 2,
+        concat_hidden: bool = False,
+        batchnorm=False,
+    ):
+        super(PreModel, self).__init__()
+        self._mask_rate = mask_rate
+        self._encoder_type = encoder_type
+        self._decoder_type = decoder_type
+        self._drop_edge_rate = drop_edge_rate
+        self._output_hidden_size = num_hidden
+        self._concat_hidden = concat_hidden
+
+        self._replace_rate = replace_rate
+        self._mask_token_rate = 1 - self._replace_rate
+
+        assert num_hidden % nhead == 0
+        assert num_hidden % nhead_out == 0
+
+        enc_num_hidden = num_hidden
+        enc_nhead = 1
+
+        dec_in_dim = num_hidden
+        dec_num_hidden = num_hidden
+
+        # Build encoder
+        self.encoder = setup_module(
+            m_type=encoder_type,
+            in_dim=in_dim,
+            out_dim=enc_num_hidden,
+            num_hidden=enc_num_hidden,
+            num_layers=num_layers,
+            activation=activation,
+            batchnorm=batchnorm,
+        )
+
+        # Build decoder for attribute prediction
+        self.decoder = setup_module(
+            m_type=decoder_type,
+            in_dim=dec_in_dim,
+            out_dim=in_dim,
+            num_hidden=dec_num_hidden,
+            num_layers=1,
+            activation=activation,
+            batchnorm=batchnorm,
+        )
+
+        self.enc_mask_token = nn.Parameter(torch.zeros(1, in_dim))
+        if concat_hidden:
+            self.encoder_to_decoder = nn.Linear(
+                dec_in_dim * num_layers, dec_in_dim, bias=False
+            )
+        else:
+            self.encoder_to_decoder = nn.Linear(dec_in_dim, dec_in_dim, bias=False)
+
+        # Setup loss function
+        self.criterion = self.setup_loss_fn(loss_fn, alpha_l)
+
+    @property
+    def output_hidden_dim(self):
+        return self._output_hidden_size
+
+    def setup_loss_fn(self, loss_fn, alpha_l):
+        if loss_fn == "mse":
+            criterion = nn.MSELoss()
+        elif loss_fn == "sce":
+            criterion = partial(sce_loss, alpha=alpha_l)
+        else:
+            raise NotImplementedError
+        return criterion
+
+    def encoding_mask_noise(self, x, mask_rate=0.3, virtual_node_index=None):
+        num_nodes = x.shape[0]
+        all_indices = torch.arange(num_nodes, device=x.device)
+
+        # Remove virtual node index from masking candidates
+        if virtual_node_index is not None:
+            all_indices = all_indices[~torch.isin(all_indices, virtual_node_index)]
+
+        perm = all_indices[torch.randperm(len(all_indices), device=x.device)]
+
+        # random masking
+        num_mask_nodes = int(mask_rate * len(perm))
+        mask_nodes = perm[:num_mask_nodes]
+        keep_nodes = perm[num_mask_nodes:]
+
+        out_x = x.clone()
+
+        if self._replace_rate > 0:
+            num_noise_nodes = int(self._replace_rate * num_mask_nodes)
+            perm_mask = torch.randperm(num_mask_nodes, device=x.device)
+            token_nodes = mask_nodes[
+                perm_mask[: int(self._mask_token_rate * num_mask_nodes)]
+            ]
+            noise_nodes = mask_nodes[
+                perm_mask[-int(self._replace_rate * num_mask_nodes) :]
+            ]
+            noise_to_be_chosen = torch.randperm(len(perm), device=x.device)[
+                :num_noise_nodes
+            ]
+            noise_to_be_chosen = all_indices[noise_to_be_chosen]
+
+            out_x[token_nodes] = 0.0
+            out_x[noise_nodes] = x[noise_to_be_chosen]
+        else:
+            token_nodes = mask_nodes
+            out_x[mask_nodes] = 0.0
+
+        out_x[token_nodes] += self.enc_mask_token
+
+        return out_x, (mask_nodes, keep_nodes)
+
+    def forward(self, batch):
+        # ---- attribute reconstruction ----
+        x, edge_index, edge_attr, virtual_node_index, batch = (
+            batch.x,
+            batch.edge_index,
+            batch.edge_attr,
+            getattr(batch, "virtual_node_index", None),
+            batch.batch,
+        )
+        loss = self.mask_attr_prediction(
+            x, edge_index, edge_attr, batch, virtual_node_index
+        )
+        return loss
+
+    def mask_attr_prediction(self, x, edge_index, edge_attr, batch, virtual_node_index):
+
+        use_x, (mask_nodes, keep_nodes) = self.encoding_mask_noise(
+            x,
+            self._mask_rate,
+            virtual_node_index,
+        )
+
+        if self._drop_edge_rate > 0:
+            use_edge_index, masked_edges = dropout_edge(
+                edge_index, self._drop_edge_rate
+            )
+            use_edge_attr = edge_attr[masked_edges]
+            use_edge_index, use_edge_attr = add_self_loops(
+                use_edge_index, use_edge_attr, fill_value="min"
+            )
+        else:
+            use_edge_index = edge_index
+            use_edge_attr = edge_attr
+
+        enc_rep, all_hidden = self.encoder(
+            use_x, use_edge_index, use_edge_attr, return_hidden=True
+        )
+        if self._concat_hidden:
+            enc_rep = torch.cat(all_hidden, dim=1)
+
+        # ---- attribute reconstruction ----
+        rep = self.encoder_to_decoder(enc_rep)
+
+        if self._decoder_type not in ("mlp", "linear"):
+            # * remask, re-mask
+            rep[mask_nodes] = 0
+
+        if self._decoder_type in ("mlp", "linear"):
+            recon = self.decoder(rep)
+        else:
+            recon = self.decoder(rep, use_edge_index, use_edge_attr)
+
+        x_init = x[mask_nodes]
+        x_rec = recon[mask_nodes]
+
+        loss = self.criterion(x_rec, x_init)
+
+        return loss
+
+    def embed(self, x, edge_index, edge_attr, batch):
+        if self._concat_hidden:
+            enc_rep, all_hidden = self.encoder(
+                x, edge_index, edge_attr, return_hidden=True
+            )
+            enc_rep = torch.cat(all_hidden, dim=1)
+        else:
+            enc_rep = self.encoder(x, edge_index, edge_attr)
+        rep = self.encoder_to_decoder(enc_rep)
+        return rep
+
+    @property
+    def enc_params(self):
+        return self.encoder.parameters()
+
+    @property
+    def dec_params(self):
+        return chain(*[self.encoder_to_decoder.parameters(), self.decoder.parameters()])

models/utils.py

@@ -0,0 +1,75 @@
+import torch
+import torch.nn as nn
+from functools import partial
+
+
+def create_activation(name):
+    if name == "relu":
+        return nn.ReLU()
+    elif name == "gelu":
+        return nn.GELU()
+    elif name == "prelu":
+        return nn.PReLU()
+    elif name is None:
+        return nn.Identity()
+    elif name == "elu":
+        return nn.ELU()
+    else:
+        raise NotImplementedError(f"{name} is not implemented.")
+
+
+def create_norm(name):
+    if name == "layernorm":
+        return nn.LayerNorm
+    elif name == "batchnorm":
+        return nn.BatchNorm1d
+    elif name == "graphnorm":
+        return partial(NormLayer, norm_type="groupnorm")
+    else:
+        return nn.Identity
+
+
+class NormLayer(nn.Module):
+    def __init__(self, hidden_dim, norm_type):
+        super().__init__()
+        if norm_type == "batchnorm":
+            self.norm = nn.BatchNorm1d(hidden_dim)
+        elif norm_type == "layernorm":
+            self.norm = nn.LayerNorm(hidden_dim)
+        elif norm_type == "graphnorm":
+            self.norm = norm_type
+            self.weight = nn.Parameter(torch.ones(hidden_dim))
+            self.bias = nn.Parameter(torch.zeros(hidden_dim))
+
+            self.mean_scale = nn.Parameter(torch.ones(hidden_dim))
+        else:
+            raise NotImplementedError
+
+    def forward(self, graph, x):
+        tensor = x
+        if self.norm is not None and type(self.norm) != str:
+            return self.norm(tensor)
+        elif self.norm is None:
+            return tensor
+
+        batch_list = graph.batch_num_nodes
+        batch_size = len(batch_list)
+        batch_list = torch.Tensor(batch_list).long().to(tensor.device)
+        batch_index = (
+            torch.arange(batch_size).to(tensor.device).repeat_interleave(batch_list)
+        )
+        batch_index = batch_index.view((-1,) + (1,) * (tensor.dim() - 1)).expand_as(
+            tensor
+        )
+        mean = torch.zeros(batch_size, *tensor.shape[1:]).to(tensor.device)
+        mean = mean.scatter_add_(0, batch_index, tensor)
+        mean = (mean.T / batch_list).T
+        mean = mean.repeat_interleave(batch_list, dim=0)
+
+        sub = tensor - mean * self.mean_scale
+
+        std = torch.zeros(batch_size, *tensor.shape[1:]).to(tensor.device)
+        std = std.scatter_add_(0, batch_index, sub.pow(2))
+        std = ((std.T / batch_list).T + 1e-6).sqrt()
+        std = std.repeat_interleave(batch_list, dim=0)
+        return self.weight * sub / std + self.bias