[Model] Training GraphSAGE with PyTorch Lightning (#2878)

* pytorch lightning initial examples

* revert most changes in dataloader to favor #2886.

* address comments

examples/pytorch/graphsage/README.md

@@ -66,3 +66,14 @@ Notably,
   in the paper.
 
 Micro F1 score reaches 0.9212 on test set.
+
+### Training with PyTorch Lightning
+
+We also provide minibatch training scripts with PyTorch Lightning in `train_lightning.py` and `train_lightning_unsupervised.py`.
+
+Requires `pytorch_lightning` and `torchmetrics`.
+
+```bash
+python3 train_lightning.py
+python3 train_lightning_unsupervised.py
+```

examples/pytorch/graphsage/model.py

+import torch as th
+import torch.nn as nn
+import torch.functional as F
+import dgl
+import dgl.nn as dglnn
+import sklearn.linear_model as lm
+import sklearn.metrics as skm
+import tqdm
+
+class SAGE(nn.Module):
+    def __init__(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout):
+        super().__init__()
+        self.init(in_feats, n_hidden, n_classes, n_layers, activation, dropout)
+
+    def init(self, in_feats, n_hidden, n_classes, n_layers, activation, dropout):
+        self.n_layers = n_layers
+        self.n_hidden = n_hidden
+        self.n_classes = n_classes
+        self.layers = nn.ModuleList()
+        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
+        for i in range(1, n_layers - 1):
+            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
+        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
+        self.dropout = nn.Dropout(dropout)
+        self.activation = activation
+
+    def forward(self, blocks, x):
+        h = x
+        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
+            h = layer(block, h)
+            if l != len(self.layers) - 1:
+                h = self.activation(h)
+                h = self.dropout(h)
+        return h
+
+    def inference(self, g, x, device, batch_size, num_workers):
+        """
+        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
+        g : the entire graph.
+        x : the input of entire node set.
+
+        The inference code is written in a fashion that it could handle any number of nodes and
+        layers.
+        """
+        # 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.
+        # TODO: can we standardize this?
+        for l, layer in enumerate(self.layers):
+            y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
+
+            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
+            dataloader = dgl.dataloading.NodeDataLoader(
+                g,
+                th.arange(g.num_nodes()).to(g.device),
+                sampler,
+                batch_size=batch_size,
+                shuffle=True,
+                drop_last=False,
+                num_workers=num_workers)
+
+            for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
+                block = blocks[0]
+
+                block = block.int().to(device)
+                h = x[input_nodes].to(device)
+                h = layer(block, h)
+                if l != len(self.layers) - 1:
+                    h = self.activation(h)
+                    h = self.dropout(h)
+
+                y[output_nodes] = h.cpu()
+
+            x = y
+        return y
+
+def compute_acc_unsupervised(emb, labels, train_nids, val_nids, test_nids):
+    """
+    Compute the accuracy of prediction given the labels.
+    """
+    emb = emb.cpu().numpy()
+    labels = labels.cpu().numpy()
+    train_nids = train_nids.cpu().numpy()
+    train_labels = labels[train_nids]
+    val_nids = val_nids.cpu().numpy()
+    val_labels = labels[val_nids]
+    test_nids = test_nids.cpu().numpy()
+    test_labels = labels[test_nids]
+
+    emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
+
+    lr = lm.LogisticRegression(multi_class='multinomial', max_iter=10000)
+    lr.fit(emb[train_nids], train_labels)
+
+    pred = lr.predict(emb)
+    f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average='micro')
+    f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average='micro')
+    return f1_micro_eval, f1_micro_test

examples/pytorch/graphsage/negative_sampler.py

+import torch as th
+import dgl
+
+class NegativeSampler(object):
+    def __init__(self, g, k, neg_share=False):
+        self.weights = g.in_degrees().float() ** 0.75
+        self.k = k
+        self.neg_share = neg_share
+
+    def __call__(self, g, eids):
+        src, _ = g.find_edges(eids)
+        n = len(src)
+        if self.neg_share and n % self.k == 0:
+            dst = self.weights.multinomial(n, replacement=True)
+            dst = dst.view(-1, 1, self.k).expand(-1, self.k, -1).flatten()
+        else:
+            dst = self.weights.multinomial(n*self.k, replacement=True)
+        src = src.repeat_interleave(self.k)
+        return src, dst

examples/pytorch/graphsage/train_lightning.py

+import dgl
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import dgl.nn.pytorch as dglnn
+import time
+import argparse
+import tqdm
+import glob
+import os
+
+from load_graph import load_reddit, inductive_split, load_ogb
+
+from torchmetrics import Accuracy
+from pytorch_lightning.callbacks import ModelCheckpoint
+from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from model import SAGE
+
+class SAGELightning(LightningModule):
+    def __init__(self,
+                 in_feats,
+                 n_hidden,
+                 n_classes,
+                 n_layers,
+                 activation,
+                 dropout,
+                 lr):
+        super().__init__()
+        self.save_hyperparameters()
+        self.module = SAGE(in_feats, n_hidden, n_classes, n_layers, activation, dropout)
+        self.lr = lr
+        # The usage of `train_acc` and `val_acc` is the recommended practice from now on as per
+        # https://torchmetrics.readthedocs.io/en/latest/pages/lightning.html
+        self.train_acc = Accuracy()
+        self.val_acc = Accuracy()
+
+    def training_step(self, batch, batch_idx):
+        input_nodes, output_nodes, mfgs = batch
+        mfgs = [mfg.int().to(device) for mfg in mfgs]
+        batch_inputs = mfgs[0].srcdata['features']
+        batch_labels = mfgs[-1].dstdata['labels']
+        batch_pred = self.module(mfgs, batch_inputs)
+        loss = F.cross_entropy(batch_pred, batch_labels)
+        self.train_acc(th.softmax(batch_pred, 1), batch_labels)
+        self.log('train_acc', self.train_acc, prog_bar=True, on_step=True, on_epoch=False)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        input_nodes, output_nodes, mfgs = batch
+        mfgs = [mfg.int().to(device) for mfg in mfgs]
+        batch_inputs = mfgs[0].srcdata['features']
+        batch_labels = mfgs[-1].dstdata['labels']
+        batch_pred = self.module(mfgs, batch_inputs)
+        self.val_acc(th.softmax(batch_pred, 1), batch_labels)
+        self.log('val_acc', self.val_acc, prog_bar=True, on_step=True, on_epoch=True, sync_dist=True)
+
+    def configure_optimizers(self):
+        optimizer = th.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+
+class DataModule(LightningDataModule):
+    def __init__(self, dataset_name, data_cpu=False, fan_out=[10, 25],
+                 device=th.device('cpu'), batch_size=1000, num_workers=4):
+        super().__init__()
+        if dataset_name == 'reddit':
+            g, n_classes = load_reddit()
+        elif dataset_name == 'ogbn-products':
+            g, n_classes = load_ogb('ogbn-products')
+        else:
+            raise ValueError('unknown dataset')
+
+        train_nid = th.nonzero(g.ndata['train_mask'], as_tuple=True)[0]
+        val_nid = th.nonzero(g.ndata['val_mask'], as_tuple=True)[0]
+        test_nid = th.nonzero(~(g.ndata['train_mask'] | g.ndata['val_mask']), as_tuple=True)[0]
+
+        sampler = dgl.dataloading.MultiLayerNeighborSampler([int(_) for _ in fan_out])
+
+        dataloader_device = th.device('cpu')
+        if not data_cpu:
+            train_nid = train_nid.to(device)
+            val_nid = val_nid.to(device)
+            test_nid = test_nid.to(device)
+            g = g.formats(['csc'])
+            g = g.to(device)
+            dataloader_device = device
+
+        self.g = g
+        self.train_nid, self.val_nid, self.test_nid = train_nid, val_nid, test_nid
+        self.sampler = sampler
+        self.device = dataloader_device
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.in_feats = g.ndata['features'].shape[1]
+        self.n_classes = n_classes
+
+    def train_dataloader(self):
+        return dgl.dataloading.NodeDataLoader(
+            self.g,
+            self.train_nid,
+            self.sampler,
+            device=self.device,
+            batch_size=self.batch_size,
+            shuffle=True,
+            drop_last=False,
+            num_workers=self.num_workers)
+
+    def val_dataloader(self):
+        return dgl.dataloading.NodeDataLoader(
+            self.g,
+            self.val_nid,
+            self.sampler,
+            device=self.device,
+            batch_size=self.batch_size,
+            shuffle=True,
+            drop_last=False,
+            num_workers=self.num_workers)
+
+
+def evaluate(model, g, val_nid, device):
+    """
+    Evaluate the model on the validation set specified by ``val_nid``.
+    g : The entire graph.
+    val_nid : the node Ids for validation.
+    device : The GPU device to evaluate on.
+    """
+    model.eval()
+    nfeat = g.ndata['features']
+    labels = g.ndata['labels']
+    with th.no_grad():
+        pred = model.module.inference(g, nfeat, device, args.batch_size, args.num_workers)
+    model.train()
+    test_acc = Accuracy()
+    return test_acc(th.softmax(pred[val_nid], -1), labels[val_nid].to(pred.device))
+
+
+if __name__ == '__main__':
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument('--gpu', type=int, default=0,
+                           help="GPU device ID. Use -1 for CPU training")
+    argparser.add_argument('--dataset', type=str, default='reddit')
+    argparser.add_argument('--num-epochs', type=int, default=20)
+    argparser.add_argument('--num-hidden', type=int, default=16)
+    argparser.add_argument('--num-layers', type=int, default=2)
+    argparser.add_argument('--fan-out', type=str, default='10,25')
+    argparser.add_argument('--batch-size', type=int, default=1000)
+    argparser.add_argument('--log-every', type=int, default=20)
+    argparser.add_argument('--eval-every', type=int, default=5)
+    argparser.add_argument('--lr', type=float, default=0.003)
+    argparser.add_argument('--dropout', type=float, default=0.5)
+    argparser.add_argument('--num-workers', type=int, default=0,
+                           help="Number of sampling processes. Use 0 for no extra process.")
+    argparser.add_argument('--inductive', action='store_true',
+                           help="Inductive learning setting")
+    argparser.add_argument('--data-cpu', action='store_true',
+                           help="By default the script puts the graph, node features and labels "
+                                "on GPU when using it to save time for data copy. This may "
+                                "be undesired if they cannot fit in GPU memory at once. "
+                                "This flag disables that.")
+    args = argparser.parse_args()
+
+    if args.gpu >= 0:
+        device = th.device('cuda:%d' % args.gpu)
+    else:
+        device = th.device('cpu')
+
+    datamodule = DataModule(
+        args.dataset, args.data_cpu, [int(_) for _ in args.fan_out.split(',')],
+        device, args.batch_size, args.num_workers)
+    model = SAGELightning(
+        datamodule.in_feats, args.num_hidden, datamodule.n_classes, args.num_layers,
+        F.relu, args.dropout, args.lr)
+
+    # Train
+    checkpoint_callback = ModelCheckpoint(monitor='val_acc', save_top_k=1)
+    trainer = Trainer(gpus=[args.gpu] if args.gpu != -1 else None,
+                      max_epochs=args.num_epochs,
+                      callbacks=[checkpoint_callback])
+    trainer.fit(model, datamodule=datamodule)
+
+    # Test
+    dirs = glob.glob('./lightning_logs/*')
+    version = max([int(os.path.split(x)[-1].split('_')[-1]) for x in dirs])
+    logdir = './lightning_logs/version_%d' % version
+    print('Evaluating model in', logdir)
+    ckpt = glob.glob(os.path.join(logdir, 'checkpoints', '*'))[0]
+
+    model = SAGELightning.load_from_checkpoint(
+        checkpoint_path=ckpt, hparams_file=os.path.join(logdir, 'hparams.yaml')).to(device)
+    test_acc = evaluate(model, datamodule.g, datamodule.test_nid, device)
+    print('Test accuracy:', test_acc)

examples/pytorch/graphsage/train_lightning_unsupervised.py

+import dgl
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import dgl.nn.pytorch as dglnn
+import dgl.function as fn
+import time
+import argparse
+import tqdm
+import glob
+import os
+
+from negative_sampler import NegativeSampler
+
+from pytorch_lightning.metrics import Accuracy
+from pytorch_lightning.callbacks import ModelCheckpoint, Callback
+from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from model import SAGE, compute_acc_unsupervised as compute_acc
+from load_graph import load_reddit, inductive_split, load_ogb
+
+class CrossEntropyLoss(nn.Module):
+    def forward(self, block_outputs, pos_graph, neg_graph):
+        with pos_graph.local_scope():
+            pos_graph.ndata['h'] = block_outputs
+            pos_graph.apply_edges(fn.u_dot_v('h', 'h', 'score'))
+            pos_score = pos_graph.edata['score']
+        with neg_graph.local_scope():
+            neg_graph.ndata['h'] = block_outputs
+            neg_graph.apply_edges(fn.u_dot_v('h', 'h', 'score'))
+            neg_score = neg_graph.edata['score']
+
+        score = th.cat([pos_score, neg_score])
+        label = th.cat([th.ones_like(pos_score), th.zeros_like(neg_score)]).long()
+        loss = F.binary_cross_entropy_with_logits(score, label.float())
+        return loss
+
+class SAGELightning(LightningModule):
+    def __init__(self,
+                 in_feats,
+                 n_hidden,
+                 n_classes,
+                 n_layers,
+                 activation,
+                 dropout,
+                 lr):
+        super().__init__()
+        self.save_hyperparameters()
+        self.module = SAGE(in_feats, n_hidden, n_classes, n_layers, activation, dropout)
+        self.lr = lr
+        self.loss_fcn = CrossEntropyLoss()
+
+    def training_step(self, batch, batch_idx):
+        input_nodes, pos_graph, neg_graph, mfgs = batch
+        mfgs = [mfg.int().to(device) for mfg in mfgs]
+        pos_graph = pos_graph.to(device)
+        neg_graph = neg_graph.to(device)
+        batch_inputs = mfgs[0].srcdata['features']
+        batch_labels = mfgs[-1].dstdata['labels']
+        batch_pred = self.module(mfgs, batch_inputs)
+        loss = self.loss_fcn(batch_pred, pos_graph, neg_graph)
+        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        input_nodes, output_nodes, mfgs = batch
+        mfgs = [mfg.int().to(device) for mfg in mfgs]
+        batch_inputs = mfgs[0].srcdata['features']
+        batch_labels = mfgs[-1].dstdata['labels']
+        batch_pred = self.module(mfgs, batch_inputs)
+        return batch_pred
+
+    def configure_optimizers(self):
+        optimizer = th.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+
+class DataModule(LightningDataModule):
+    def __init__(self, dataset_name, data_cpu=False, fan_out=[10, 25],
+                 device=th.device('cpu'), batch_size=1000, num_workers=4):
+        super().__init__()
+        if dataset_name == 'reddit':
+            g, n_classes = load_reddit()
+            n_edges = g.num_edges()
+            reverse_eids = th.cat([
+                th.arange(n_edges // 2, n_edges),
+                th.arange(0, n_edges // 2)])
+        elif dataset_name == 'ogbn-products':
+            g, n_classes = load_ogb('ogbn-products')
+            n_edges = g.num_edges()
+            # The reverse edge of edge 0 in OGB products dataset is 1.
+            # The reverse edge of edge 2 is 3.  So on so forth.
+            reverse_eids = th.arange(n_edges) ^ 1
+        else:
+            raise ValueError('unknown dataset')
+
+        train_nid = th.nonzero(g.ndata['train_mask'], as_tuple=True)[0]
+        val_nid = th.nonzero(g.ndata['val_mask'], as_tuple=True)[0]
+        test_nid = th.nonzero(~(g.ndata['train_mask'] | g.ndata['val_mask']), as_tuple=True)[0]
+
+        sampler = dgl.dataloading.MultiLayerNeighborSampler([int(_) for _ in fan_out])
+
+        dataloader_device = th.device('cpu')
+        if not data_cpu:
+            train_nid = train_nid.to(device)
+            val_nid = val_nid.to(device)
+            test_nid = test_nid.to(device)
+            g = g.formats(['csc'])
+            g = g.to(device)
+            dataloader_device = device
+
+        self.g = g
+        self.train_nid, self.val_nid, self.test_nid = train_nid, val_nid, test_nid
+        self.sampler = sampler
+        self.device = dataloader_device
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.in_feats = g.ndata['features'].shape[1]
+        self.n_classes = n_classes
+        self.reverse_eids = reverse_eids
+
+    def train_dataloader(self):
+        return dgl.dataloading.EdgeDataLoader(
+            self.g,
+            np.arange(self.g.num_edges()),
+            self.sampler,
+            exclude='reverse_id',
+            reverse_eids=self.reverse_eids,
+            negative_sampler=NegativeSampler(self.g, args.num_negs, args.neg_share),
+            device=self.device,
+            batch_size=self.batch_size,
+            shuffle=True,
+            drop_last=False,
+            num_workers=self.num_workers)
+
+    def val_dataloader(self):
+        # Note that the validation data loader is a NodeDataLoader
+        # as we want to evaluate all the node embeddings.
+        return dgl.dataloading.NodeDataLoader(
+            self.g,
+            np.arange(self.g.num_nodes()),
+            self.sampler,
+            device=self.device,
+            batch_size=self.batch_size,
+            shuffle=False,
+            drop_last=False,
+            num_workers=self.num_workers)
+
+
+class UnsupervisedClassification(Callback):
+    def on_validation_epoch_start(self, trainer, pl_module):
+        self.val_outputs = []
+
+    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+        self.val_outputs.append(outputs)
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        node_emb = th.cat(self.val_outputs, 0)
+        g = trainer.datamodule.g
+        labels = g.ndata['labels']
+        f1_micro, f1_macro = compute_acc(
+            node_emb, labels, trainer.datamodule.train_nid,
+            trainer.datamodule.val_nid, trainer.datamodule.test_nid)
+        pl_module.log('val_f1_micro', f1_micro)
+
+if __name__ == '__main__':
+    argparser = argparse.ArgumentParser("multi-gpu training")
+    argparser.add_argument("--gpu", type=int, default=0)
+    argparser.add_argument('--dataset', type=str, default='reddit')
+    argparser.add_argument('--num-epochs', type=int, default=20)
+    argparser.add_argument('--num-hidden', type=int, default=16)
+    argparser.add_argument('--num-layers', type=int, default=2)
+    argparser.add_argument('--num-negs', type=int, default=1)
+    argparser.add_argument('--neg-share', default=False, action='store_true',
+                           help="sharing neg nodes for positive nodes")
+    argparser.add_argument('--fan-out', type=str, default='10,25')
+    argparser.add_argument('--batch-size', type=int, default=10000)
+    argparser.add_argument('--log-every', type=int, default=20)
+    argparser.add_argument('--eval-every', type=int, default=1000)
+    argparser.add_argument('--lr', type=float, default=0.003)
+    argparser.add_argument('--dropout', type=float, default=0.5)
+    argparser.add_argument('--num-workers', type=int, default=0,
+                           help="Number of sampling processes. Use 0 for no extra process.")
+    args = argparser.parse_args()
+
+    if args.gpu >= 0:
+        device = th.device('cuda:%d' % args.gpu)
+    else:
+        device = th.device('cpu')
+
+    datamodule = DataModule(
+        args.dataset, True, [int(_) for _ in args.fan_out.split(',')],
+        device, args.batch_size, args.num_workers)
+    model = SAGELightning(
+        datamodule.in_feats, args.num_hidden, datamodule.n_classes, args.num_layers,
+        F.relu, args.dropout, args.lr)
+
+    # Train
+    unsupervised_callback = UnsupervisedClassification()
+    checkpoint_callback = ModelCheckpoint(monitor='val_f1_micro', save_top_k=1)
+    trainer = Trainer(gpus=[args.gpu] if args.gpu != -1 else None,
+                      max_epochs=args.num_epochs,
+                      val_check_interval=1000,
+                      callbacks=[checkpoint_callback, unsupervised_callback],
+                      num_sanity_val_steps=0)
+    trainer.fit(model, datamodule=datamodule)

examples/pytorch/graphsage/train_sampling.py

@@ -9,79 +9,9 @@ import time
 import argparse
 import tqdm
 
+from model import SAGE
 from load_graph import load_reddit, inductive_split, load_ogb
 
-class SAGE(nn.Module):
-    def __init__(self,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
-        super().__init__()
-        self.n_layers = n_layers
-        self.n_hidden = n_hidden
-        self.n_classes = n_classes
-        self.layers = nn.ModuleList()
-        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
-        for i in range(1, n_layers - 1):
-            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
-        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
-        self.dropout = nn.Dropout(dropout)
-        self.activation = activation
-
-    def forward(self, blocks, x):
-        h = x
-        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
-            h = layer(block, h)
-            if l != len(self.layers) - 1:
-                h = self.activation(h)
-                h = self.dropout(h)
-        return h
-
-    def inference(self, g, x, device):
-        """
-        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
-        g : the entire graph.
-        x : the input of entire node set.
-
-        The inference code is written in a fashion that it could handle any number of nodes and
-        layers.
-        """
-        # 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.
-        # TODO: can we standardize this?
-        for l, layer in enumerate(self.layers):
-            y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
-
-            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
-            dataloader = dgl.dataloading.NodeDataLoader(
-                g,
-                th.arange(g.num_nodes()),
-                sampler,
-                batch_size=args.batch_size,
-                shuffle=True,
-                drop_last=False,
-                num_workers=args.num_workers)
-
-            for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
-                block = blocks[0]
-
-                block = block.int().to(device)
-                h = x[input_nodes].to(device)
-                h = layer(block, h)
-                if l != len(self.layers) - 1:
-                    h = self.activation(h)
-                    h = self.dropout(h)
-
-                y[output_nodes] = h.cpu()
-
-            x = y
-        return y
-
 def compute_acc(pred, labels):
     """
     Compute the accuracy of prediction given the labels.
@@ -100,7 +30,7 @@ def evaluate(model, g, nfeat, labels, val_nid, device):
     """
     model.eval()
     with th.no_grad():
-        pred = model.inference(g, nfeat, device)
+        pred = model.inference(g, nfeat, device, args.batch_size, args.num_workers)
     model.train()
     return compute_acc(pred[val_nid], labels[val_nid].to(pred.device))
 
@@ -192,7 +122,7 @@ def run(args, device, data):
     print('Avg epoch time: {}'.format(avg / (epoch - 4)))
 
 if __name__ == '__main__':
-    argparser = argparse.ArgumentParser("multi-gpu training")
+    argparser = argparse.ArgumentParser()
     argparser.add_argument('--gpu', type=int, default=0,
                            help="GPU device ID. Use -1 for CPU training")
     argparser.add_argument('--dataset', type=str, default='reddit')

examples/pytorch/graphsage/train_sampling_multi_gpu.py

@@ -12,80 +12,10 @@ import argparse
 from torch.nn.parallel import DistributedDataParallel
 import tqdm
 
+from model import SAGE
 from utils import thread_wrapped_func
 from load_graph import load_reddit, inductive_split
 
-class SAGE(nn.Module):
-    def __init__(self,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
-        super().__init__()
-        self.n_layers = n_layers
-        self.n_hidden = n_hidden
-        self.n_classes = n_classes
-        self.layers = nn.ModuleList()
-        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
-        for i in range(1, n_layers - 1):
-            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
-        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
-        self.dropout = nn.Dropout(dropout)
-        self.activation = activation
-
-    def forward(self, blocks, x):
-        h = x
-        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
-            h = layer(block, h)
-            if l != len(self.layers) - 1:
-                h = self.activation(h)
-                h = self.dropout(h)
-        return h
-
-    def inference(self, g, x, device):
-        """
-        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
-        g : the entire graph.
-        x : the input of entire node set.
-
-        The inference code is written in a fashion that it could handle any number of nodes and
-        layers.
-        """
-        # 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.
-        # TODO: can we standardize this?
-        for l, layer in enumerate(self.layers):
-            y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
-
-            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
-            dataloader = dgl.dataloading.NodeDataLoader(
-                g,
-                th.arange(g.num_nodes()),
-                sampler,
-                batch_size=args.batch_size,
-                shuffle=True,
-                drop_last=False,
-                num_workers=args.num_workers)
-
-            for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
-                block = blocks[0]
-
-                block = block.int().to(device)
-                h = x[input_nodes].to(device)
-                h = layer(block, h)
-                if l != len(self.layers) - 1:
-                    h = self.activation(h)
-                    h = self.dropout(h)
-
-                y[output_nodes] = h.cpu()
-
-            x = y
-        return y
-
 def compute_acc(pred, labels):
     """
     Compute the accuracy of prediction given the labels.
@@ -103,7 +33,7 @@ def evaluate(model, g, nfeat, labels, val_nid, device):
     """
     model.eval()
     with th.no_grad():
-        pred = model.inference(g, nfeat, device)
+        pred = model.inference(g, nfeat, device, args.batch_size, args.num_workers)
     model.train()
     return compute_acc(pred[val_nid], labels[val_nid])
 

examples/pytorch/graphsage/train_sampling_unsupervised.py

@@ -12,97 +12,10 @@ import argparse
 from dgl.data import RedditDataset
 from torch.nn.parallel import DistributedDataParallel
 import tqdm
-import sklearn.linear_model as lm
-import sklearn.metrics as skm
 
 from utils import thread_wrapped_func
-
-class NegativeSampler(object):
-    def __init__(self, g, k, neg_share=False):
-        self.weights = g.in_degrees().float() ** 0.75
-        self.k = k
-        self.neg_share = neg_share
-
-    def __call__(self, g, eids):
-        src, _ = g.find_edges(eids)
-        n = len(src)
-        if self.neg_share and n % self.k == 0:
-            dst = self.weights.multinomial(n, replacement=True)
-            dst = dst.view(-1, 1, self.k).expand(-1, self.k, -1).flatten()
-        else:
-            dst = self.weights.multinomial(n*self.k, replacement=True)
-        src = src.repeat_interleave(self.k)
-        return src, dst
-
-class SAGE(nn.Module):
-    def __init__(self,
-                 in_feats,
-                 n_hidden,
-                 n_classes,
-                 n_layers,
-                 activation,
-                 dropout):
-        super().__init__()
-        self.n_layers = n_layers
-        self.n_hidden = n_hidden
-        self.n_classes = n_classes
-        self.layers = nn.ModuleList()
-        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
-        for i in range(1, n_layers - 1):
-            self.layers.append(dglnn.SAGEConv(n_hidden, n_hidden, 'mean'))
-        self.layers.append(dglnn.SAGEConv(n_hidden, n_classes, 'mean'))
-        self.dropout = nn.Dropout(dropout)
-        self.activation = activation
-
-    def forward(self, blocks, x):
-        h = x
-        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
-            h = layer(block, h)
-            if l != len(self.layers) - 1:
-                h = self.activation(h)
-                h = self.dropout(h)
-        return h
-
-    def inference(self, g, x, device):
-        """
-        Inference with the GraphSAGE model on full neighbors (i.e. without neighbor sampling).
-        g : the entire graph.
-        x : the input of entire node set.
-
-        The inference code is written in a fashion that it could handle any number of nodes and
-        layers.
-        """
-        # 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.
-        # TODO: can we standardize this?
-        for l, layer in enumerate(self.layers):
-            y = th.zeros(g.num_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)
-
-            sampler = dgl.dataloading.MultiLayerFullNeighborSampler(1)
-            dataloader = dgl.dataloading.NodeDataLoader(
-                g,
-                th.arange(g.num_nodes()),
-                sampler,
-                batch_size=args.batch_size,
-                shuffle=True,
-                drop_last=False,
-                num_workers=args.num_workers)
-
-            for input_nodes, output_nodes, blocks in tqdm.tqdm(dataloader):
-                block = blocks[0].to(device)
-
-                h = x[input_nodes].to(device)
-                h = layer(block, h)
-                if l != len(self.layers) - 1:
-                    h = self.activation(h)
-                    h = self.dropout(h)
-
-                y[output_nodes] = h.cpu()
-
-            x = y
-        return y
+from model import SAGE, compute_acc_unsupervised as compute_acc
+from negative_sampler import NegativeSampler
 
 class CrossEntropyLoss(nn.Module):
     def forward(self, block_outputs, pos_graph, neg_graph):
@@ -120,29 +33,6 @@ class CrossEntropyLoss(nn.Module):
         loss = F.binary_cross_entropy_with_logits(score, label.float())
         return loss
 
-def compute_acc(emb, labels, train_nids, val_nids, test_nids):
-    """
-    Compute the accuracy of prediction given the labels.
-    """
-    emb = emb.cpu().numpy()
-    labels = labels.cpu().numpy()
-    train_nids = train_nids.cpu().numpy()
-    train_labels = labels[train_nids]
-    val_nids = val_nids.cpu().numpy()
-    val_labels = labels[val_nids]
-    test_nids = test_nids.cpu().numpy()
-    test_labels = labels[test_nids]
-
-    emb = (emb - emb.mean(0, keepdims=True)) / emb.std(0, keepdims=True)
-
-    lr = lm.LogisticRegression(multi_class='multinomial', max_iter=10000)
-    lr.fit(emb[train_nids], train_labels)
-
-    pred = lr.predict(emb)
-    f1_micro_eval = skm.f1_score(val_labels, pred[val_nids], average='micro')
-    f1_micro_test = skm.f1_score(test_labels, pred[test_nids], average='micro')
-    return f1_micro_eval, f1_micro_test
-
 def evaluate(model, g, nfeat, labels, train_nids, val_nids, test_nids, device):
     """
     Evaluate the model on the validation set specified by ``val_mask``.
@@ -156,10 +46,10 @@ def evaluate(model, g, nfeat, labels, train_nids, val_nids, test_nids, device):
     with th.no_grad():
         # single gpu
         if isinstance(model, SAGE):
-            pred = model.inference(g, nfeat, device)
+            pred = model.inference(g, nfeat, device, args.batch_size, args.num_workers)
         # multi gpu
         else:
-            pred = model.module.inference(g, nfeat, device)
+            pred = model.module.inference(g, nfeat, device, args.batch_size, args.num_workers)
     model.train()
     return compute_acc(pred, labels, train_nids, val_nids, test_nids)
 
@@ -203,7 +93,7 @@ def run(proc_id, n_gpus, args, devices, data):
         reverse_eids=th.cat([
             th.arange(n_edges // 2, n_edges),
             th.arange(0, n_edges // 2)]),
-        negative_sampler=NegativeSampler(g, args.num_negs),
+        negative_sampler=NegativeSampler(g, args.num_negs, args.neg_share),
         batch_size=args.batch_size,
         shuffle=True,
         drop_last=False,
@@ -312,7 +202,7 @@ def main(args, devices):
 if __name__ == '__main__':
     argparser = argparse.ArgumentParser("multi-gpu training")
     argparser.add_argument("--gpu", type=str, default='0',
-                           help="GPU, can be a list of gpus for multi-gpu trianing,"
+                           help="GPU, can be a list of gpus for multi-gpu training,"
                                 " e.g., 0,1,2,3; -1 for CPU")
     argparser.add_argument('--num-epochs', type=int, default=20)
     argparser.add_argument('--num-hidden', type=int, default=16)

python/dgl/dataloading/pytorch/__init__.py

@@ -15,6 +15,7 @@ class _ScalarDataBatcherIter:
         self.index = 0
         self.drop_last = drop_last
 
+    # Make this an iterator for PyTorch Lightning compatibility
     def __iter__(self):
         return self
 
@@ -67,6 +68,10 @@ class _ScalarDataBatcher(th.utils.data.IterableDataset):
 
         return _ScalarDataBatcherIter(dataset, self.batch_size, self.drop_last)
 
+    def __len__(self):
+        return (self.dataset.shape[0] + (0 if self.drop_last else self.batch_size - 1)) // \
+            self.batch_size
+
 def _remove_kwargs_dist(kwargs):
     if 'num_workers' in kwargs:
         del kwargs['num_workers']
@@ -217,6 +222,7 @@ class _NodeDataLoaderIter:
         self.node_dataloader = node_dataloader
         self.iter_ = iter(node_dataloader.dataloader)
 
+    # Make this an iterator for PyTorch Lightning compatibility
     def __iter__(self):
         return self
 
@@ -234,6 +240,7 @@ class _EdgeDataLoaderIter:
         self.edge_dataloader = edge_dataloader
         self.iter_ = iter(edge_dataloader.dataloader)
 
+    # Make this an iterator for PyTorch Lightning compatibility
     def __iter__(self):
         return self