train_sampling_reddit.py

# -*- coding: utf-8 -*-
"""
Created on Tue Jul 14 19:51:00 2020

@author: Ming Jin

Minibatch training (Algorithm 2) on Reddit dataset 
with the concept of bipartite graph

Build version:
    
    + PyTorch 1.1.0
    
    + DGL 0.4.3.post2
"""

import dgl
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
from dgl.data import RedditDataset
import tqdm

from SageConv import SAGEConv_bigraph


class NeighborSampler(object):
    """
    Line 1 to 7 of the algorithm 2 in the paper
    
    g: DGLGraph
    fanouts: sampling neighbors, len(fanouts) is the number of returned blocks
    
    seeds: minibatched nodes that we we want to generate representations for
    
    Return:
        blocks: a set of bipartite graphs in the order of [block K, ..., block 0]
                block 0: Seeds(dst) and 1-hop neighbors(src)
                block K: K-hop neighbors(dst) and (K+1)-hop neighbors(src)
    
    Example: 
        If we set batch_size=1000, then the size of the seeds will be 1000;
        If fanouts = '10,25', then we sample 10 neighbors per seed and will 
        have 1-hop neighbors; After this, we sample 25 neighbors per 1-hop
        neighbors to get 2-hop neighbors; Seeds and 1-hop neighbors form the
        block 0, 1-hop neighbors and 2-hop neighbors form the block 1;
        
    """
    def __init__(self, g, fanouts):
        
        self.g = g
        self.fanouts = fanouts

    def sample_blocks(self, seeds):
        
        # Based on the idea of bipartite graph, sampling will be performed
        # from the LHS to RHS, i.e. Seeds to their neighbors
        # Notice: Outmost neighbors, i.e. blocks[0].srcdata[dgl.NID], includes 
        # all of the nodes that will be needed for k-hops aggregation of the 
        # bathc_size seeds
        
        seeds = torch.LongTensor(np.asarray(seeds))
        blocks = []
        for fanout in self.fanouts: 
            if fanout is None:
                # This will be used during the inference
                # return the subgraph contains ALL 1-hop neighbors
                frontier = dgl.transform.in_subgraph(self.g, seeds)
            else:
                # sample_neighbors() samples 'fanout' neighbors of 'seeds' on 'g'
                # TODO: the meaning of replace=True needs to check
                frontier = dgl.sampling.sample_neighbors(self.g, seeds, fanout, replace=True) 
            # to_block() converts 'frontier' to a bipartite graph with dst as 'seeds'
            # 'include_dst_in_src=True' means to include DST nodes in SRC nodes
            # Since DST nodes are included in SRC nodes, we could fetch
            # the DST node features from the SRC nodes features (this is why I wrote
            # such a notice at the begining of this method).
            block = dgl.to_block(frontier, seeds, include_dst_in_src=True)
            # assign the SRC of the current block as the DST of next block
            seeds = block.srcdata[dgl.NID]
            # store blocks with the stack structure
            # if there are two layers (K=2), then blocks = [block 1, block 0]
            # the block_id represents loops
            blocks.insert(0, block)
            
        return blocks


class GraphSAGE(nn.Module):
    """
    Minibatch-based SAGE network
    
    Line 8 to 16 of the algorithm 2 in the paper
    
    Parameters:
        blocks: a set of bipartite graphs in the order of [block K, ..., block 0]
        x: feature matrix with shape [block0.DSTNODES, NumFeats]
    """
    
    def __init__(self,
                 in_feats,
                 n_hidden,
                 n_classes,
                 n_layers,
                 activation,
                 aggregator_type,
                 dropout = 0.5):
        
        super(GraphSAGE, self).__init__()
        self.n_hidden = n_hidden
        self.n_classes = n_classes
        self.layers = nn.ModuleList()
        self.dropout = nn.Dropout(p=dropout)
        # first layer
        self.layers.append(SAGEConv_bigraph(in_feats, n_hidden, aggregator_type, activation))
        # hidden layers
        for i in range(1, n_layers - 1):
            self.layers.append(SAGEConv_bigraph(n_hidden, n_hidden, aggregator_type, activation))
        # last layer
        self.layers.append(SAGEConv_bigraph(n_hidden, n_classes, aggregator_type, None))
        
    def forward(self, blocks, x):
        
        # Based on the idea of bipartite graph, aggregation will be performed
        # layer-wise from the RHS to LHS, i.e. Outmost neighbors to seeds
        # input x with the shape [block0.DSTNODES, NumFeats]
        # return h with the shape [seeds(batch_size), n_classes]
        # Example:
        #       h维度变化的例子：
        #       loop1: h.shape: torch.Size([105727, 602])
        #       loop2: h.shape: torch.Size([9644, 16])
        #       return h.shape: torch.Size([1000, 41])
        
        h = x
        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
            # h: [BLOCK.DST, BLOCK.SRC]
            h_dst = h[:block.number_of_dst_nodes()]
            h = layer(block, (h, h_dst))
            # we don't need activation and dropout for the last layer
            # activation has been included inside of the SAGEConv
            if l != len(self.layers) - 1:
                h = self.dropout(h)
        return h
    
    def inference(self, g, x, batch_size, device):

        # Inference with the model on full neighbors (i.e. without neighbor sampling)
        # g: the entire graph
        # x: feature matrix of the entire graph
        # batch_size: we can't feed all node to a layer at once so we need batch_size
        
        # P.S. During inference with sampling, multi-layer blocks are very inefficient
        # because lots of computations in the first few layers are repeated.
        # Therefore, we compute the representation of all nodes layer-by-layer
        # The nodes on each layer are of course splitted in batches.
        
        # Predict node representations on graph g layer-by-layer based on the 
        # aggregation order we performed on forward()
        for l, layer in enumerate(self.layers):
            
            # we use 'y' to store the k-th layer predictions of nodes on 'g'
            y = torch.zeros(g.number_of_nodes(), 
                          self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
            
            # go through nodes in 'g' batch-by-batch with 'dataloader'
            # we don't need to assign how many neighbors we want to sample cuz
            # WE WANT ALL OF NEIGHBORS OF EACH NODE
            sampler = NeighborSampler(g, [None])
            dataloader = DataLoader(
                dataset=torch.arange(g.number_of_nodes()),
                batch_size=batch_size,
                collate_fn=sampler.sample_blocks,
                shuffle=True,
                drop_last=False)

            for blocks in tqdm.tqdm(dataloader):
                # Notice: len(blocks) = 1
                block = blocks[0]
                # Example:
                #   input_nodes.shape: torch.Size([138408])
                #   output_nodes.shape: torch.Size([1000])
                input_nodes = block.srcdata[dgl.NID]
                output_nodes = block.dstdata[dgl.NID]
                # h: feature matrix of BLOCK.SRC, which is 1-hop neighbors of the output_nodes
                h = x[input_nodes].to(device)
                # h_dst: feature matrix of BLOCK.DST, which equals to the number of batch_size
                h_dst = h[:block.number_of_dst_nodes()]
                # feed into the corresponding layer to get k-th layer predictions
                # of this batch of nodes, i.e. output_nodes
                h = layer(block, (h, h_dst))
                if l != len(self.layers) - 1:
                    h = self.dropout(h)
                # inner-looping to collect k-th layer predictions of ALL nodes in 'g'
                y[output_nodes] = h.cpu()
            # k-th layer predictions will act as the features of (k+1)-th layer
            x = y
        # outer-looping goes through layers to generate the final node embeddings of 'g'
        # Example:
        #   return y.shape: torch.Size([232965, 41])
        #   232965 is NumNodes，41 is n_classes
        return y

def compute_acc(pred, labels):
    """
    Return the accuracy of prediction given the labels. 
    """
    return (torch.argmax(pred, dim=1) == labels).float().sum() / len(pred)

def evaluate(model, g, inputs, labels, val_nid, batch_size, device):
    """
    Evaluate the model on the validation set specified by 'val_nid'
    g : The entire graph.
    inputs : The features of all the nodes, i.e. g.ndata['features']
    labels : The labels of all the nodes, i.e. torch.LongTensor(data.labels)
    val_nid : the node Ids for validation
    batch_size : Number of nodes to compute at the same time
    device : The GPU device to evaluate on
    """
    model.eval()
    with torch.no_grad():
        pred = model.inference(g, inputs, batch_size, device)
    model.train()
    return compute_acc(pred[val_nid], labels[val_nid])

def load_subtensor(g, seeds, input_nodes, device):
    """
    Copys features and labels of a set of nodes onto GPU.
    """
    batch_inputs = g.ndata['features'][input_nodes].to(device)
    batch_labels = g.ndata['labels'][seeds].to(device)
    return batch_inputs, batch_labels


'''
Entry point
'''

# Parameters
gpu = 0
aggregation = 'mean'
num_epochs = 20
num_hidden = 256
num_layers = 2  # n_layers >= 2
fan_out = '10,25'
batch_size = 512
log_every = 20  # frequency of printout
eval_every = 10
lr = 0.0003
dropout = 0.5

if gpu >= 0:
    device = torch.device('cuda:%d' % gpu)
else:
    device = torch.device('cpu')

# Reddit dataset

# NumNodes: 232965
# NumEdges: 114848857
# NumFeats: 602
# NumClasses: 41
# NumTrainingSamples: 153431
# NumValidationSamples: 23831
# NumTestSamples: 55703

data = RedditDataset(self_loop=True)
train_mask = data.train_mask
val_mask = data.val_mask
test_mask = data.test_mask
features = torch.Tensor(data.features)
labels = torch.LongTensor(data.labels)
in_feats = features.shape[1]
n_classes = data.num_labels

# Construct graph
g = dgl.graph(data.graph.all_edges())
g.ndata['features'] = features
g.ndata['labels'] = labels

# get different node IDs
# Examples:
#   train_nid.shape: 153431
#   val_nid.shape: 23831
#   test_nid.shape: 55703   
train_nid = torch.LongTensor(np.nonzero(train_mask)[0])
val_nid = torch.LongTensor(np.nonzero(val_mask)[0])
test_nid = torch.LongTensor(np.nonzero(test_mask)[0])

# Create training sampler
sampler = NeighborSampler(g, [int(fanout) for fanout in fan_out.split(',')])

# Create PyTorch DataLoader for constructing blocks
# collate_fn <-- sampler，which could sample k-hops neighbors for seeds in a batch
dataloader = DataLoader(
    dataset=train_nid.numpy(),
    batch_size=batch_size,
    collate_fn=sampler.sample_blocks,
    shuffle=True,
    drop_last=False)

# Define model and optimizer
model = GraphSAGE(in_feats, 
                  num_hidden, 
                  n_classes, 
                  num_layers, 
                  nn.ReLU(), 
                  aggregation, 
                  dropout)
model = model.to(device)
loss_fcn = nn.CrossEntropyLoss()
loss_fcn = loss_fcn.to(device)
optimizer = optim.Adam(model.parameters(), lr=lr)

# Training loop
avg = 0
iter_tput = []
for epoch in range(num_epochs):
    tic = time.time()
    # loop over nodes in the training graph
    # return a list of blocks for each batch of seed nodes
    # len(blocks): should equals to num_layers
    # If K=2 then:
    #    blocks[0] is block 1 in sample_blocks(), represents 1-hop & 2-hop bigraph
    #    blocks[1] is block 0 in sample_blocks(), represents seeds & 1-hop bigraph
    # Example (if batch_size=1000, fanout='10,25'):
    #    blocks[0].srcdata[dgl.NID].shape: torch.Size([106421])
    #    blocks[0].dstdata[dgl.NID].shape: torch.Size([9670])
    #    blocks[1].srcdata[dgl.NID].shape: torch.Size([9670])
    #    blocks[1].dstdata[dgl.NID].shape: torch.Size([1000])
    for step, blocks in enumerate(dataloader):
        
        tic_step = time.time()
        
        # As we mentioned before, input_nodes contains all nodes we need
        # This could also be found in line 3 and 5 in algorithm 2 
        input_nodes = blocks[0].srcdata[dgl.NID]
        # seeds is a batch of nodes that we want to calculate their embeddings
        seeds = blocks[-1].dstdata[dgl.NID]
        # Example: 
        #   input_nodes.shape: torch.Size([106421])
        #   seeds.shape: torch.Size([1000])
        
        # Load the input features as well as output labels
        # Example:
        #   batch_inputs.shape: torch.Size([106421, 602])
        #   batch_labels.shape: torch.Size([1000])
        batch_inputs, batch_labels = load_subtensor(g, seeds, input_nodes, device)
        
        # Compute loss and prediction
        # batch_pred with the shape [batch_size, n_classes]
        batch_pred = model(blocks, batch_inputs)
         
        optimizer.zero_grad()
        
        # Notice: CrossEntropyLoss contains softmax() and negative log loss
        loss = loss_fcn(batch_pred, batch_labels)
        loss.backward()
        optimizer.step()

        iter_tput.append(len(seeds) / (time.time() - tic_step))
        if step % log_every == 0:
            # training accuracy
            acc = compute_acc(batch_pred, batch_labels)
            # GPU mem usage
            gpu_mem_alloc = torch.cuda.max_memory_allocated() / 1000000 if torch.cuda.is_available() else 0
            # printout
            print('Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MiB'.format(
                epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc))
    
    # Epoch time calculation
    toc = time.time()
    print('Epoch Time(s): {:.4f}'.format(toc - tic))
    if epoch >= 5:
        avg += toc - tic
    # Evaluate on both val_set and test_set
    if epoch % eval_every == 0 and epoch != 0:
        eval_acc = evaluate(model, g, g.ndata['features'], labels, val_nid, batch_size, device)
        print('Eval Acc {:.4f}'.format(eval_acc))
        test_acc = evaluate(model, g, g.ndata['features'], labels, test_nid, batch_size, device)
        print('Test Acc: {:.4f}'.format(test_acc))

test_acc = evaluate(model, g, g.ndata['features'], labels, test_nid, batch_size, device)
print('Test Acc: {:.4f}'.format(test_acc))
print('Avg epoch time: {}'.format(avg / (epoch - 4)))