add base to hetero setting

test/performance/heterogeneous/dgl/han_base/base.py

@@ -0,0 +1,108 @@
+import torch
+from sklearn.metrics import f1_score
+
+from utils import load_data, EarlyStopping
+
+def score(logits, labels):
+    _, indices = torch.max(logits, dim=1)
+    prediction = indices.long().cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    accuracy = (prediction == labels).sum() / len(prediction)
+    micro_f1 = f1_score(labels, prediction, average='micro')
+    macro_f1 = f1_score(labels, prediction, average='macro')
+
+    return accuracy, micro_f1, macro_f1
+
+def evaluate(model, g, features, labels, mask, loss_func):
+    model.eval()
+    with torch.no_grad():
+        logits = model(g, features)
+    loss = loss_func(logits[mask], labels[mask])
+    accuracy, micro_f1, macro_f1 = score(logits[mask], labels[mask])
+
+    return loss, accuracy, micro_f1, macro_f1
+
+def main(args):
+    # If args['hetero'] is True, g would be a heterogeneous graph.
+    # Otherwise, it will be a list of homogeneous graphs.
+    g, features, labels, num_classes, train_idx, val_idx, test_idx, train_mask, \
+    val_mask, test_mask = load_data(args['dataset'])
+
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = train_mask.bool()
+        val_mask = val_mask.bool()
+        test_mask = test_mask.bool()
+
+    features = features.to(args['device'])
+    labels = labels.to(args['device'])
+    train_mask = train_mask.to(args['device'])
+    val_mask = val_mask.to(args['device'])
+    test_mask = test_mask.to(args['device'])
+
+    if args['hetero']:
+        from model_hetero import HAN
+        model = HAN(meta_paths=[['pa', 'ap'], ['pf', 'fp']],
+                    in_size=features.shape[1],
+                    hidden_size=args['hidden_units'],
+                    out_size=num_classes,
+                    num_heads=args['num_heads'],
+                    dropout=args['dropout']).to(args['device'])
+        g = g.to(args['device'])
+    else:
+        from model import HAN
+        model = HAN(num_meta_paths=len(g),
+                    in_size=features.shape[1],
+                    hidden_size=args['hidden_units'],
+                    out_size=num_classes,
+                    num_heads=args['num_heads'],
+                    dropout=args['dropout']).to(args['device'])
+        g = [graph.to(args['device']) for graph in g]
+
+    stopper = EarlyStopping(patience=args['patience'])
+    loss_fcn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'],
+                                 weight_decay=args['weight_decay'])
+
+    for epoch in range(args['num_epochs']):
+        model.train()
+        logits = model(g, features)
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        train_acc, train_micro_f1, train_macro_f1 = score(logits[train_mask], labels[train_mask])
+        val_loss, val_acc, val_micro_f1, val_macro_f1 = evaluate(model, g, features, labels, val_mask, loss_fcn)
+        early_stop = stopper.step(val_loss.data.item(), val_acc, model)
+
+        print('Epoch {:d} | Train Loss {:.4f} | Train Micro f1 {:.4f} | Train Macro f1 {:.4f} | '
+              'Val Loss {:.4f} | Val Micro f1 {:.4f} | Val Macro f1 {:.4f}'.format(
+            epoch + 1, loss.item(), train_micro_f1, train_macro_f1, val_loss.item(), val_micro_f1, val_macro_f1))
+
+        if early_stop:
+            break
+
+    stopper.load_checkpoint(model)
+    test_loss, test_acc, test_micro_f1, test_macro_f1 = evaluate(model, g, features, labels, test_mask, loss_fcn)
+    print('Test loss {:.4f} | Test Micro f1 {:.4f} | Test Macro f1 {:.4f}'.format(
+        test_loss.item(), test_micro_f1, test_macro_f1))
+
+if __name__ == '__main__':
+    import argparse
+
+    from utils import setup
+
+    parser = argparse.ArgumentParser('HAN')
+    parser.add_argument('-s', '--seed', type=int, default=1,
+                        help='Random seed')
+    parser.add_argument('-ld', '--log-dir', type=str, default='results',
+                        help='Dir for saving training results')
+    parser.add_argument('--hetero', action='store_true',
+                        help='Use metapath coalescing with DGL\'s own dataset')
+    args = parser.parse_args().__dict__
+
+    args = setup(args)
+
+    main(args)

test/performance/heterogeneous/dgl/han_base/model.py

@@ -0,0 +1,83 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+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
+    ---------
+    num_meta_paths : number of homogeneous graphs generated from the metapaths.
+    in_size : input feature dimension
+    out_size : output feature dimension
+    layer_num_heads : number of attention heads
+    dropout : Dropout probability
+
+    Inputs
+    ------
+    g : list[DGLGraph]
+        List of graphs
+    h : tensor
+        Input features
+
+    Outputs
+    -------
+    tensor
+        The output feature
+    """
+    def __init__(self, num_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(num_meta_paths):
+            self.gat_layers.append(GATConv(in_size, out_size, layer_num_heads,
+                                           dropout, dropout, activation=F.elu))
+        self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
+        self.num_meta_paths = num_meta_paths
+
+    def forward(self, gs, h):
+        semantic_embeddings = []
+
+        for i, g in enumerate(gs):
+            semantic_embeddings.append(self.gat_layers[i](g, h).flatten(1))
+        semantic_embeddings = torch.stack(semantic_embeddings, dim=1)                  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)                            # (N, D * K)
+
+class HAN(nn.Module):
+    def __init__(self, num_meta_paths, in_size, hidden_size, out_size, num_heads, dropout):
+        super(HAN, self).__init__()
+
+        self.layers = nn.ModuleList()
+        self.layers.append(HANLayer(num_meta_paths, in_size, hidden_size, num_heads[0], dropout))
+        for l in range(1, len(num_heads)):
+            self.layers.append(HANLayer(num_meta_paths, hidden_size * num_heads[l-1],
+                                        hidden_size, num_heads[l], dropout))
+        self.predict = nn.Linear(hidden_size * num_heads[-1], out_size)
+
+    def forward(self, g, h):
+        for gnn in self.layers:
+            h = gnn(g, h)
+
+        return self.predict(h)

test/performance/heterogeneous/dgl/han_base/model_hetero.py

@@ -0,0 +1,105 @@
+"""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))
+        self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
+        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))
+        semantic_embeddings = torch.stack(semantic_embeddings, dim=1)                  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)                            # (N, D * K)
+
+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))
+        self.predict = nn.Linear(hidden_size * num_heads[-1], out_size)
+
+    def forward(self, g, h):
+        for gnn in self.layers:
+            h = gnn(g, h)
+
+        return self.predict(h)

test/performance/heterogeneous/dgl/han_base/train_sampling.py

@@ -0,0 +1,271 @@
+# -*- coding: utf-8 -*-
+"""
+HAN mini-batch training by RandomWalkSampler.
+note: This demo use RandomWalkSampler to sample neighbors, it's hard to get all neighbors when valid or test,
+so we sampled twice as many neighbors during val/test than training.
+"""
+import dgl
+import numpy
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dgl.nn.pytorch import GATConv
+
+from dgl.sampling import RandomWalkNeighborSampler
+from sklearn.metrics import f1_score
+from torch.utils.data import DataLoader
+
+from model_hetero import SemanticAttention
+from utils import EarlyStopping, set_random_seed
+
+
+class HANLayer(torch.nn.Module):
+    """
+    HAN layer.
+
+    Arguments
+    ---------
+    num_metapath : number of metapath based sub-graph
+    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, num_metapath, 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(num_metapath):
+            self.gat_layers.append(GATConv(in_size, out_size, layer_num_heads,
+                                           dropout, dropout, activation=F.elu,
+                                           allow_zero_in_degree=True))
+        self.semantic_attention = SemanticAttention(in_size=out_size * layer_num_heads)
+        self.num_metapath = num_metapath
+
+    def forward(self, block_list, h_list):
+        semantic_embeddings = []
+
+        for i, block in enumerate(block_list):
+            semantic_embeddings.append(self.gat_layers[i](block, h_list[i]).flatten(1))
+        semantic_embeddings = torch.stack(semantic_embeddings, dim=1)  # (N, M, D * K)
+
+        return self.semantic_attention(semantic_embeddings)  # (N, D * K)
+
+
+class HAN(nn.Module):
+    def __init__(self, num_metapath, in_size, hidden_size, out_size, num_heads, dropout):
+        super(HAN, self).__init__()
+
+        self.layers = nn.ModuleList()
+        self.layers.append(HANLayer(num_metapath, in_size, hidden_size, num_heads[0], dropout))
+        for l in range(1, len(num_heads)):
+            self.layers.append(HANLayer(num_metapath, hidden_size * num_heads[l - 1],
+                                        hidden_size, num_heads[l], dropout))
+        self.predict = nn.Linear(hidden_size * num_heads[-1], out_size)
+
+    def forward(self, g, h):
+        for gnn in self.layers:
+            h = gnn(g, h)
+
+        return self.predict(h)
+
+
+class HANSampler(object):
+    def __init__(self, g, metapath_list, num_neighbors):
+        self.sampler_list = []
+        for metapath in metapath_list:
+            # note: random walk may get same route(same edge), which will be removed in the sampled graph.
+            # So the sampled graph's edges may be less than num_random_walks(num_neighbors).
+            self.sampler_list.append(RandomWalkNeighborSampler(G=g,
+                                                               num_traversals=1,
+                                                               termination_prob=0,
+                                                               num_random_walks=num_neighbors,
+                                                               num_neighbors=num_neighbors,
+                                                               metapath=metapath))
+
+    def sample_blocks(self, seeds):
+        block_list = []
+        for sampler in self.sampler_list:
+            frontier = sampler(seeds)
+            # add self loop
+            frontier = dgl.remove_self_loop(frontier)
+            frontier.add_edges(torch.tensor(seeds), torch.tensor(seeds))
+            block = dgl.to_block(frontier, seeds)
+            block_list.append(block)
+
+        return seeds, block_list
+
+
+def score(logits, labels):
+    _, indices = torch.max(logits, dim=1)
+    prediction = indices.long().cpu().numpy()
+    labels = labels.cpu().numpy()
+
+    accuracy = (prediction == labels).sum() / len(prediction)
+    micro_f1 = f1_score(labels, prediction, average='micro')
+    macro_f1 = f1_score(labels, prediction, average='macro')
+
+    return accuracy, micro_f1, macro_f1
+
+
+def evaluate(model, g, metapath_list, num_neighbors, features, labels, val_nid, loss_fcn, batch_size):
+    model.eval()
+
+    han_valid_sampler = HANSampler(g, metapath_list, num_neighbors=num_neighbors * 2)
+    dataloader = DataLoader(
+        dataset=val_nid,
+        batch_size=batch_size,
+        collate_fn=han_valid_sampler.sample_blocks,
+        shuffle=False,
+        drop_last=False,
+        num_workers=4)
+    correct = total = 0
+    prediction_list = []
+    labels_list = []
+    with torch.no_grad():
+        for step, (seeds, blocks) in enumerate(dataloader):
+            h_list = load_subtensors(blocks, features)
+            blocks = [block.to(args['device']) for block in blocks]
+            hs = [h.to(args['device']) for h in h_list]
+
+            logits = model(blocks, hs)
+            loss = loss_fcn(logits, labels[numpy.asarray(seeds)].to(args['device']))
+            # get each predict label
+            _, indices = torch.max(logits, dim=1)
+            prediction = indices.long().cpu().numpy()
+            labels_batch = labels[numpy.asarray(seeds)].cpu().numpy()
+
+            prediction_list.append(prediction)
+            labels_list.append(labels_batch)
+
+            correct += (prediction == labels_batch).sum()
+            total += prediction.shape[0]
+
+    total_prediction = numpy.concatenate(prediction_list)
+    total_labels = numpy.concatenate(labels_list)
+    micro_f1 = f1_score(total_labels, total_prediction, average='micro')
+    macro_f1 = f1_score(total_labels, total_prediction, average='macro')
+    accuracy = correct / total
+
+    return loss, accuracy, micro_f1, macro_f1
+
+
+def load_subtensors(blocks, features):
+    h_list = []
+    for block in blocks:
+        input_nodes = block.srcdata[dgl.NID]
+        h_list.append(features[input_nodes])
+    return h_list
+
+
+def main(args):
+    # acm data
+    if args['dataset'] == 'ACMRaw':
+        from utils import load_data
+        g, features, labels, n_classes, train_nid, val_nid, test_nid, train_mask, \
+        val_mask, test_mask = load_data('ACMRaw')
+        metapath_list = [['pa', 'ap'], ['pf', 'fp']]
+    else:
+        raise NotImplementedError('Unsupported dataset {}'.format(args['dataset']))
+
+    # Is it need to set different neighbors numbers for different meta-path based graph?
+    num_neighbors = args['num_neighbors']
+    han_sampler = HANSampler(g, metapath_list, num_neighbors)
+    # Create PyTorch DataLoader for constructing blocks
+    dataloader = DataLoader(
+        dataset=train_nid,
+        batch_size=args['batch_size'],
+        collate_fn=han_sampler.sample_blocks,
+        shuffle=True,
+        drop_last=False,
+        num_workers=4)
+
+    model = HAN(num_metapath=len(metapath_list),
+                in_size=features.shape[1],
+                hidden_size=args['hidden_units'],
+                out_size=n_classes,
+                num_heads=args['num_heads'],
+                dropout=args['dropout']).to(args['device'])
+
+    total_params = sum(p.numel() for p in model.parameters())
+    print("total_params: {:d}".format(total_params))
+    total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("total trainable params: {:d}".format(total_trainable_params))
+
+    stopper = EarlyStopping(patience=args['patience'])
+    loss_fn = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'],
+                                 weight_decay=args['weight_decay'])
+
+    for epoch in range(args['num_epochs']):
+        model.train()
+        for step, (seeds, blocks) in enumerate(dataloader):
+            h_list = load_subtensors(blocks, features)
+            blocks = [block.to(args['device']) for block in blocks]
+            hs = [h.to(args['device']) for h in h_list]
+
+            logits = model(blocks, hs)
+            loss = loss_fn(logits, labels[numpy.asarray(seeds)].to(args['device']))
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # print info in each batch
+            train_acc, train_micro_f1, train_macro_f1 = score(logits, labels[numpy.asarray(seeds)])
+            print(
+                "Epoch {:d} | loss: {:.4f} | train_acc: {:.4f} | train_micro_f1: {:.4f} | train_macro_f1: {:.4f}".format(
+                    epoch + 1, loss, train_acc, train_micro_f1, train_macro_f1
+                ))
+        val_loss, val_acc, val_micro_f1, val_macro_f1 = evaluate(model, g, metapath_list, num_neighbors, features,
+                                                                 labels, val_nid, loss_fn, args['batch_size'])
+        early_stop = stopper.step(val_loss.data.item(), val_acc, model)
+
+        print('Epoch {:d} | Val loss {:.4f} | Val Accuracy {:.4f} | Val Micro f1 {:.4f} | Val Macro f1 {:.4f}'.format(
+            epoch + 1, val_loss.item(), val_acc, val_micro_f1, val_macro_f1))
+
+        if early_stop:
+            break
+
+    stopper.load_checkpoint(model)
+    test_loss, test_acc, test_micro_f1, test_macro_f1 = evaluate(model, g, metapath_list, num_neighbors, features,
+                                                                 labels, test_nid, loss_fn, args['batch_size'])
+    print('Test loss {:.4f} | Test Accuracy {:.4f} | Test Micro f1 {:.4f} | Test Macro f1 {:.4f}'.format(
+        test_loss.item(), test_acc, test_micro_f1, test_macro_f1))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('mini-batch HAN')
+    parser.add_argument('-s', '--seed', type=int, default=1,
+                        help='Random seed')
+    parser.add_argument('--batch_size', type=int, default=32)
+    parser.add_argument('--num_neighbors', type=int, default=20)
+    parser.add_argument('--lr', type=float, default=0.001)
+    parser.add_argument('--num_heads', type=list, default=[8])
+    parser.add_argument('--hidden_units', type=int, default=8)
+    parser.add_argument('--dropout', type=float, default=0.6)
+    parser.add_argument('--weight_decay', type=float, default=0.001)
+    parser.add_argument('--num_epochs', type=int, default=100)
+    parser.add_argument('--patience', type=int, default=10)
+    parser.add_argument('--dataset', type=str, default='ACMRaw')
+    parser.add_argument('--device', type=str, default='cuda:0')
+
+    args = parser.parse_args().__dict__
+    # set_random_seed(args['seed'])
+
+    main(args)

test/performance/heterogeneous/dgl/han_base/utils.py

@@ -0,0 +1,275 @@
+import sys
+import datetime
+import dgl
+import errno
+import numpy as np
+import os
+import pickle
+import random
+import torch
+
+from dgl.data.utils import download, get_download_dir, _get_dgl_url
+import os.path as osp
+from pprint import pprint
+from scipy import sparse
+from scipy import io as sio
+
+def set_random_seed(seed=0):
+    """Set random seed.
+    Parameters
+    ----------
+    seed : int
+        Random seed to use
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+
+def mkdir_p(path, log=True):
+    """Create a directory for the specified path.
+    Parameters
+    ----------
+    path : str
+        Path name
+    log : bool
+        Whether to print result for directory creation
+    """
+    try:
+        os.makedirs(path)
+        if log:
+            print('Created directory {}'.format(path))
+    except OSError as exc:
+        if exc.errno == errno.EEXIST and os.path.isdir(path) and log:
+            print('Directory {} already exists.'.format(path))
+        else:
+            raise
+
+def get_date_postfix():
+    """Get a date based postfix for directory name.
+    Returns
+    -------
+    post_fix : str
+    """
+    dt = datetime.datetime.now()
+    post_fix = '{}_{:02d}-{:02d}-{:02d}'.format(
+        dt.date(), dt.hour, dt.minute, dt.second)
+
+    return post_fix
+
+def setup_log_dir(args, sampling=False):
+    """Name and create directory for logging.
+    Parameters
+    ----------
+    args : dict
+        Configuration
+    Returns
+    -------
+    log_dir : str
+        Path for logging directory
+    sampling : bool
+        Whether we are using sampling based training
+    """
+    date_postfix = get_date_postfix()
+    log_dir = os.path.join(
+        args['log_dir'],
+        '{}_{}'.format(args['dataset'], date_postfix))
+
+    if sampling:
+        log_dir = log_dir + '_sampling'
+
+    mkdir_p(log_dir)
+    return log_dir
+
+# The configuration below is from the paper.
+default_configure = {
+    'lr': 0.005,             # Learning rate
+    'num_heads': [8],        # Number of attention heads for node-level attention
+    'hidden_units': 8,
+    'dropout': 0.6,
+    'weight_decay': 0.001,
+    'num_epochs': 200,
+    'patience': 100
+}
+
+sampling_configure = {
+    'batch_size': 20
+}
+
+def setup(args):
+    args.update(default_configure)
+    set_random_seed(args['seed'])
+    args['dataset'] = 'ACMRaw' if args['hetero'] else 'ACM'
+    args['device'] = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    args['log_dir'] = setup_log_dir(args)
+    return args
+
+def setup_for_sampling(args):
+    args.update(default_configure)
+    args.update(sampling_configure)
+    set_random_seed()
+    args['device'] = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    args['log_dir'] = setup_log_dir(args, sampling=True)
+    return args
+
+def get_binary_mask(total_size, indices):
+    mask = torch.zeros(total_size)
+    mask[indices] = 1
+    return mask.byte()
+
+def load_acm(remove_self_loop):
+    filename = 'ACM3025.pkl'
+    url = 'dataset/' + filename
+    data_path = get_download_dir() + '/' + filename
+    if osp.exists(data_path):
+        print(f'Using existing file {filename}', file=sys.stderr)
+    else:
+        download(_get_dgl_url(url), path=data_path)
+
+    with open(data_path, 'rb') as f:
+        data = pickle.load(f)
+
+    labels, features = torch.from_numpy(data['label'].todense()).long(), \
+                       torch.from_numpy(data['feature'].todense()).float()
+    num_classes = labels.shape[1]
+    labels = labels.nonzero()[:, 1]
+
+    if remove_self_loop:
+        num_nodes = data['label'].shape[0]
+        data['PAP'] = sparse.csr_matrix(data['PAP'] - np.eye(num_nodes))
+        data['PLP'] = sparse.csr_matrix(data['PLP'] - np.eye(num_nodes))
+
+    # Adjacency matrices for meta path based neighbors
+    # (Mufei): I verified both of them are binary adjacency matrices with self loops
+    author_g = dgl.from_scipy(data['PAP'])
+    subject_g = dgl.from_scipy(data['PLP'])
+    gs = [author_g, subject_g]
+
+    train_idx = torch.from_numpy(data['train_idx']).long().squeeze(0)
+    val_idx = torch.from_numpy(data['val_idx']).long().squeeze(0)
+    test_idx = torch.from_numpy(data['test_idx']).long().squeeze(0)
+
+    num_nodes = author_g.number_of_nodes()
+    train_mask = get_binary_mask(num_nodes, train_idx)
+    val_mask = get_binary_mask(num_nodes, val_idx)
+    test_mask = get_binary_mask(num_nodes, test_idx)
+
+    print('dataset loaded')
+    pprint({
+        'dataset': 'ACM',
+        'train': train_mask.sum().item() / num_nodes,
+        'val': val_mask.sum().item() / num_nodes,
+        'test': test_mask.sum().item() / num_nodes
+    })
+
+    return gs, features, labels, num_classes, train_idx, val_idx, test_idx, \
+           train_mask, val_mask, test_mask
+
+def load_acm_raw(remove_self_loop):
+    assert not remove_self_loop
+    filename = 'ACM.mat'
+    url = 'dataset/' + filename
+    data_path = get_download_dir() + '/' + filename
+    if osp.exists(data_path):
+        print(f'Using existing file {filename}', file=sys.stderr)
+    else:
+        download(_get_dgl_url(url), path=data_path)
+
+    data = sio.loadmat(data_path)
+    p_vs_l = data['PvsL']       # paper-field?
+    p_vs_a = data['PvsA']       # paper-author
+    p_vs_t = data['PvsT']       # paper-term, bag of words
+    p_vs_c = data['PvsC']       # paper-conference, labels come from that
+
+    # We assign
+    # (1) KDD papers as class 0 (data mining),
+    # (2) SIGMOD and VLDB papers as class 1 (database),
+    # (3) SIGCOMM and MOBICOMM papers as class 2 (communication)
+    conf_ids = [0, 1, 9, 10, 13]
+    label_ids = [0, 1, 2, 2, 1]
+
+    p_vs_c_filter = p_vs_c[:, conf_ids]
+    p_selected = (p_vs_c_filter.sum(1) != 0).A1.nonzero()[0]
+    p_vs_l = p_vs_l[p_selected]
+    p_vs_a = p_vs_a[p_selected]
+    p_vs_t = p_vs_t[p_selected]
+    p_vs_c = p_vs_c[p_selected]
+
+    hg = dgl.heterograph({
+        ('paper', 'pa', 'author'): p_vs_a.nonzero(),
+        ('author', 'ap', 'paper'): p_vs_a.transpose().nonzero(),
+        ('paper', 'pf', 'field'): p_vs_l.nonzero(),
+        ('field', 'fp', 'paper'): p_vs_l.transpose().nonzero()
+    })
+
+    features = torch.FloatTensor(p_vs_t.toarray())
+
+    pc_p, pc_c = p_vs_c.nonzero()
+    labels = np.zeros(len(p_selected), dtype=np.int64)
+    for conf_id, label_id in zip(conf_ids, label_ids):
+        labels[pc_p[pc_c == conf_id]] = label_id
+    labels = torch.LongTensor(labels)
+
+    num_classes = 3
+
+    float_mask = np.zeros(len(pc_p))
+    for conf_id in conf_ids:
+        pc_c_mask = (pc_c == conf_id)
+        float_mask[pc_c_mask] = np.random.permutation(np.linspace(0, 1, pc_c_mask.sum()))
+    train_idx = np.where(float_mask <= 0.2)[0]
+    val_idx = np.where((float_mask > 0.2) & (float_mask <= 0.3))[0]
+    test_idx = np.where(float_mask > 0.3)[0]
+
+    num_nodes = hg.number_of_nodes('paper')
+    train_mask = get_binary_mask(num_nodes, train_idx)
+    val_mask = get_binary_mask(num_nodes, val_idx)
+    test_mask = get_binary_mask(num_nodes, test_idx)
+
+    return hg, features, labels, num_classes, train_idx, val_idx, test_idx, \
+            train_mask, val_mask, test_mask
+
+def load_data(dataset, remove_self_loop=False):
+    if dataset == 'ACM':
+        return load_acm(remove_self_loop)
+    elif dataset == 'ACMRaw':
+        return load_acm_raw(remove_self_loop)
+    else:
+        return NotImplementedError('Unsupported dataset {}'.format(dataset))
+
+class EarlyStopping(object):
+    def __init__(self, patience=10):
+        dt = datetime.datetime.now()
+        self.filename = 'early_stop_{}_{:02d}-{:02d}-{:02d}.pth'.format(
+            dt.date(), dt.hour, dt.minute, dt.second)
+        self.patience = patience
+        self.counter = 0
+        self.best_acc = None
+        self.best_loss = None
+        self.early_stop = False
+
+    def step(self, loss, acc, model):
+        if self.best_loss is None:
+            self.best_acc = acc
+            self.best_loss = loss
+            self.save_checkpoint(model)
+        elif (loss > self.best_loss) and (acc < self.best_acc):
+            self.counter += 1
+            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            if (loss <= self.best_loss) and (acc >= self.best_acc):
+                self.save_checkpoint(model)
+            self.best_loss = np.min((loss, self.best_loss))
+            self.best_acc = np.max((acc, self.best_acc))
+            self.counter = 0
+        return self.early_stop
+
+    def save_checkpoint(self, model):
+        """Saves model when validation loss decreases."""
+        torch.save(model.state_dict(), self.filename)
+
+    def load_checkpoint(self, model):
+        """Load the latest checkpoint."""
+        model.load_state_dict(torch.load(self.filename))

test/performance/node_classification/dgl/han_main.py → test/performance/heterogeneous/dgl/han_main.py

@@ -29,6 +29,7 @@ from scipy import sparse
 from scipy import io as sio 
 #from model import *
 from sklearn.metrics import f1_score
+import autogl.module.model
 
 class EarlyStopping(object):
     def __init__(self, patience=10):