Class GraphMAE

code/model.py:204–277

    #     return pred

class GraphMAE(nn.Module):

    EPS = 1e-15

    def __init__(self, in_dim, emb_dim, num_layer, kernel='gcn', drop_ratio=0,
                 act='relu', norm=None, concat=False, mask_ratio=0.5, replace_ratio=0):
        super().__init__()

        self.emb_dim = emb_dim if not concat else num_layer * emb_dim
        self.mask_ratio = mask_ratio
        self.replace_ratio = replace_ratio

        self.encoder = Encoder(in_dim, emb_dim, num_layer, kernel, drop_ratio, act, norm, 
                               concat=False, last_act=True, aggr='mean')
        self.decoder = Encoder(emb_dim, in_dim, 1, kernel, drop_ratio, act, norm=None, 
                               concat=False, last_act=False, aggr='mean')

        self.encoder_mask_token = nn.Parameter(torch.zeros(1, in_dim))
        self.encoder_to_decoder = nn.Linear(self.emb_dim, emb_dim, bias=False)


    def forward(self, x, edge_index, edge_weigt=None, batch=None):
        # Mask
        mask_x, mask_nodes = self.encding_mask(x)
        h = self.encoder(mask_x, edge_index, edge_weigt, batch)
        h = self.encoder_to_decoder(h)

        # Re-mask
        h[mask_nodes] = 0
        recon = self.decoder(h, edge_index, edge_weigt, batch)

        return x[mask_nodes], recon[mask_nodes]


    def encding_mask(self, x):

        num_nodes = x.shape[0]

        # random masking
        perm = torch.randperm(num_nodes, device=x.device)
        num_mask_nodes = int(self.mask_ratio * num_nodes)
        mask_nodes = perm[: num_mask_nodes]
        # keep_nodes = perm[num_mask_nodes: ]

        if self.replace_ratio > 0:
            num_noise_nodes = int(self.replace_ratio * num_mask_nodes)
            perm_mask = torch.randperm(num_mask_nodes, device=x.device)
            token_nodes = mask_nodes[perm_mask[: int((1 - self.replace_ratio) * num_mask_nodes)]]
            noise_nodes = mask_nodes[perm_mask[-int(self.replace_ratio * num_mask_nodes):]]
            noise_to_be_chosen = torch.randperm(num_nodes, device=x.device)[:num_noise_nodes]

            out_x = x.clone()
            out_x[token_nodes] = 0.0
            out_x[noise_nodes] = x[noise_to_be_chosen]
        else:
            out_x = x.clone()
            token_nodes = mask_nodes
            out_x[mask_nodes] = 0.0