Class GraphInfoMax

code/model.py:152–202

    

class GraphInfoMax(nn.Module):

    EPS = 1e-15

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

        self.emd_dim = emb_dim

        self.encoder = Encoder(in_dim, emb_dim, num_layer, kernel, drop_ratio, act, norm, concat=False, last_act=False)
        # self.pred_head = LinearPred(emb_dim, emb_dim, 2, 2)

        self.weight = nn.Parameter(torch.empty(emb_dim, emb_dim))
        # just try
        uniform(self.emd_dim, self.weight)
        # self.weight = nn.Parameter(torch.empty(768, 768))
        # uniform(768, self.weight)
        

    def forward(self, x, edge_index, edge_weigt=None, batch=None):

        pos_h = self.encoder(x, edge_index, edge_weigt, batch)

        x_cor = infomax_corruption(x, batch)
        neg_h = self.encoder(x_cor, edge_index, edge_weigt, batch)

        summary = torch.sigmoid(pos_h.mean(dim=0))

        return pos_h, neg_h, summary
    
    def discriminate(self, h, summary):

        summary = summary.t() if summary.dim() > 1 else summary
        value = torch.matmul(h, torch.matmul(self.weight, summary))
        return torch.sigmoid(value)
    
    def get_loss(self, pos_h, neg_h, summary):

        pos_loss = -torch.log(self.discriminate(pos_h, summary) + self.EPS).mean()
        neg_loss = -torch.log(1 - self.discriminate(neg_h, summary) + self.EPS).mean()

        return pos_loss + neg_loss
    

    # def predict(self, x, edge_index, edge_weigt=None, batch=None):
        
    #     h, _, _ = self.forward(x, edge_index, edge_weigt, batch)
    #     pred = F.softmax(self.pred_head(h), dim=-1)

    #     return pred

class GraphMAE(nn.Module):