model_hetero.py

"""This model shows an example of using dgl.metapath_reachable_graph on the original heterogeneous
graph.
Because the original HAN implementation only gives the preprocessed homogeneous graph, this model
could not reproduce the result in HAN as they did not provide the preprocessing code, and we
constructed another dataset from ACM with a different set of papers, connections, features and
labels.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

import dgl
from dgl.nn.pytorch import GATConv

class SemanticAttention(nn.Module):
    def __init__(self, in_size, hidden_size=128):
        super(SemanticAttention, self).__init__()

        self.project = nn.Sequential(
            nn.Linear(in_size, hidden_size),
            nn.Tanh(),
            nn.Linear(hidden_size, 1, bias=False)
        )

    def forward(self, z):
        w = self.project(z).mean(0)                    # (M, 1)
        beta = torch.softmax(w, dim=0)                 # (M, 1)
        beta = beta.expand((z.shape[0],) + beta.shape) # (N, M, 1)

        return (beta * z).sum(1)                       # (N, D * K)

class HANLayer(nn.Module):
    """
    HAN layer.
    Arguments
    ---------
    meta_paths : list of metapaths, each as a list of edge types
    in_size : input feature dimension
    out_size : output feature dimension
    layer_num_heads : number of attention heads
    dropout : Dropout probability
    Inputs
    ------
    g : DGLHeteroGraph
        The heterogeneous graph
    h : tensor
        Input features
    Outputs
    -------
    tensor
        The output feature
    """
    def __init__(self, meta_paths, in_size, out_size, layer_num_heads, dropout):
        super(HANLayer, self).__init__()

        # One GAT layer for each meta path based adjacency matrix
        self.gat_layers = nn.ModuleList()
        for i in range(len(meta_paths)):
            self.gat_layers.append(GATConv(in_size, out_size, layer_num_heads,
                                           dropout, dropout, activation=F.elu,
                                           allow_zero_in_degree=True))
        if len(meta_paths) > 1:
            self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
        else:
            self.semantic_attention = None
        self.meta_paths = list(tuple(meta_path) for meta_path in meta_paths)

        self._cached_graph = None
        self._cached_coalesced_graph = {}

    def forward(self, g, h):
        semantic_embeddings = []

        if self._cached_graph is None or self._cached_graph is not g:
            self._cached_graph = g
            self._cached_coalesced_graph.clear()
            for meta_path in self.meta_paths:
                self._cached_coalesced_graph[meta_path] = dgl.metapath_reachable_graph(
                        g, meta_path)

        for i, meta_path in enumerate(self.meta_paths):
            new_g = self._cached_coalesced_graph[meta_path]
            semantic_embeddings.append(self.gat_layers[i](new_g, h).flatten(1))

        if self.semantic_attention is not None:
            semantic_embeddings = torch.stack(semantic_embeddings, dim=1)                  # (N, M, D * K)
            return self.semantic_attention(semantic_embeddings)                            # (N, D * K)
        else:
            return semantic_embeddings[-1]

class HAN(nn.Module):
    def __init__(self, meta_paths, in_size, hidden_size, out_size, num_heads, dropout):
        super(HAN, self).__init__()

        self.layers = nn.ModuleList()
        self.layers.append(HANLayer(meta_paths, in_size, hidden_size, num_heads[0], dropout))
        for l in range(1, len(num_heads)):
            self.layers.append(HANLayer(meta_paths, hidden_size * num_heads[l-1],
                                        hidden_size, num_heads[l], dropout))


    def forward(self, g, h):
        for gnn in self.layers:
            h = gnn(g, h)

        return h

class SS_HAN(nn.Module):
    def __init__(self, muti_meta_paths, in_size, hidden_size, out_size, num_heads, dropout):
        super(SS_HAN, self).__init__()

        self.han = nn.ModuleList()
        for meta_paths in muti_meta_paths:
            self.han.append(HAN(meta_paths, in_size, hidden_size, out_size, num_heads, dropout))

    def forward(self, g, h):
        embeddings = []
        for i, han in enumerate(self.han):
            embeddings.append(han(g, h[i]))

        return embeddings

    def calculate_loss(self, embeddings, pos_edge_index, neg_edge_index):
        edge_index = torch.cat([pos_edge_index, neg_edge_index], dim=-1)
        z = torch.cat(embeddings, dim=0)
        logits = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=-1)
        return logits

class Classifier(nn.Module):
    def __init__(self, in_size, hidden_size, out_size):
        super(Classifier, self).__init__()

        self.hidden_layer = nn.Linear(in_size, hidden_size)
        self.output_layer = nn.Linear(hidden_size, out_size)

    def forward(self, x):

        x = self.hidden_layer(x)
        x = F.relu(x)
        x = F.dropout(x)
        x = self.output_layer(x)

        return x