Delete models

models/__init__.py

@@ -1,55 +0,0 @@
-"""
-(c) Adaptation of the code from https://github.com/THUDM/GraphMAE
-"""
-
-from .edcoder import PreModel
-
-
-def build_model(args):
-    num_heads = args.num_heads
-    num_out_heads = args.num_out_heads
-    num_hidden = args.num_hidden
-    num_layers = args.num_layers
-    residual = args.residual
-    attn_drop = args.attn_drop
-    in_drop = args.in_drop
-    norm = args.norm
-    negative_slope = args.negative_slope
-    encoder_type = args.encoder
-    decoder_type = args.decoder
-    mask_rate = args.mask_rate
-    drop_edge_rate = args.drop_edge_rate
-    replace_rate = args.replace_rate
-    batchnorm = args.batchnorm
-
-    activation = args.activation
-    loss_fn = args.loss_fn
-    alpha_l = args.alpha_l
-    concat_hidden = args.concat_hidden
-    num_features = args.num_features
-    num_edge_features = args.num_edge_features
-
-    model = PreModel(
-        in_dim=int(num_features),
-        edge_in_dim=int(num_edge_features),
-        num_hidden=int(num_hidden),
-        num_layers=num_layers,
-        nhead=num_heads,
-        nhead_out=num_out_heads,
-        activation=activation,
-        feat_drop=in_drop,
-        attn_drop=attn_drop,
-        negative_slope=negative_slope,
-        residual=residual,
-        encoder_type=encoder_type,
-        decoder_type=decoder_type,
-        mask_rate=mask_rate,
-        norm=norm,
-        loss_fn=loss_fn,
-        drop_edge_rate=drop_edge_rate,
-        replace_rate=replace_rate,
-        alpha_l=alpha_l,
-        concat_hidden=concat_hidden,
-        batchnorm=batchnorm,
-    )
-    return model

models/acm_gin.py

@@ -1,207 +0,0 @@
-"""
-(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 models.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

@@ -1,260 +0,0 @@
-"""
-(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 .acm_gin import ACM_GIN_model
-from torch_geometric.utils import dropout_edge
-from torch_geometric.utils import add_self_loops
-
-
-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

@@ -1,75 +0,0 @@
-import torch.nn as nn
-import torch
-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