train_sampling.py 9.58 KB
Newer Older
1
2
3
4
5
6
7
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 torch.multiprocessing as mp
8
from torch.utils.data import DataLoader
9
10
11
12
13
14
15
16
import dgl.function as fn
import dgl.nn.pytorch as dglnn
import time
import argparse
from _thread import start_new_thread
from functools import wraps
from dgl.data import RedditDataset
import tqdm
17
import traceback
18

19
from load_graph import load_reddit, load_ogb, inductive_split
20

21
22
23
24
25
26
27
28
29
30
31
32
33
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()
34
        self.layers.append(dglnn.SAGEConv(in_feats, n_hidden, 'mean'))
35
        for i in range(1, n_layers - 1):
36
37
38
39
            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
40
41
42

    def forward(self, blocks, x):
        h = x
43
        for l, (layer, block) in enumerate(zip(self.layers, blocks)):
44
45
46
47
            # We need to first copy the representation of nodes on the RHS from the
            # appropriate nodes on the LHS.
            # Note that the shape of h is (num_nodes_LHS, D) and the shape of h_dst
            # would be (num_nodes_RHS, D)
48
            h_dst = h[:block.number_of_dst_nodes()]
49
50
51
            # Then we compute the updated representation on the RHS.
            # The shape of h now becomes (num_nodes_RHS, D)
            h = layer(block, (h, h_dst))
52
53
54
            if l != len(self.layers) - 1:
                h = self.activation(h)
                h = self.dropout(h)
55
56
57
58
59
60
61
        return h

    def inference(self, g, x, batch_size, 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.
62

63
64
65
66
67
68
69
70
71
72
73
        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.number_of_nodes(), self.n_hidden if l != len(self.layers) - 1 else self.n_classes)

74
75
76
77
78
79
80
81
82
83
84
85
            sampler = dgl.sampling.MultiLayerNeighborSampler([None])
            dataloader = dgl.sampling.NodeDataLoader(
                g,
                th.arange(g.number_of_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]
86

87
                h = x[input_nodes].to(device)
88
                h_dst = h[:block.number_of_dst_nodes()]
89
                h = layer(block, (h, h_dst))
90
91
92
                if l != len(self.layers) - 1:
                    h = self.activation(h)
                    h = self.dropout(h)
93

94
                y[output_nodes] = h.cpu()
95
96
97
98

            x = y
        return y

99
100
101
102
103
def prepare_mp(g):
    """
    Explicitly materialize the CSR, CSC and COO representation of the given graph
    so that they could be shared via copy-on-write to sampler workers and GPU
    trainers.
104

105
106
107
108
109
110
    This is a workaround before full shared memory support on heterogeneous graphs.
    """
    g.in_degree(0)
    g.out_degree(0)
    g.find_edges([0])

111
112
113
114
def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
115
    labels = labels.long()
116
117
    return (th.argmax(pred, dim=1) == labels).float().sum() / len(pred)

118
def evaluate(model, g, inputs, labels, val_nid, batch_size, device):
119
    """
120
    Evaluate the model on the validation set specified by ``val_nid``.
121
122
123
    g : The entire graph.
    inputs : The features of all the nodes.
    labels : The labels of all the nodes.
124
    val_nid : the node Ids for validation.
125
126
127
128
129
130
131
    batch_size : Number of nodes to compute at the same time.
    device : The GPU device to evaluate on.
    """
    model.eval()
    with th.no_grad():
        pred = model.inference(g, inputs, batch_size, device)
    model.train()
132
    return compute_acc(pred[val_nid], labels[val_nid])
133

134
def load_subtensor(g, seeds, input_nodes, device):
135
136
137
    """
    Copys features and labels of a set of nodes onto GPU.
    """
138
    batch_inputs = g.ndata['features'][input_nodes].to(device)
139
    batch_labels = g.ndata['labels'][seeds].to(device)
140
141
142
    return batch_inputs, batch_labels

#### Entry point
143
def run(args, device, data):
144
    # Unpack data
145
146
147
148
    in_feats, n_classes, train_g, val_g, test_g = data
    train_nid = th.nonzero(train_g.ndata['train_mask'], as_tuple=True)[0]
    val_nid = th.nonzero(val_g.ndata['val_mask'], as_tuple=True)[0]
    test_nid = th.nonzero(~(test_g.ndata['train_mask'] | test_g.ndata['val_mask']), as_tuple=True)[0]
149

150
    # Create PyTorch DataLoader for constructing blocks
151
152
153
    sampler = dgl.sampling.MultiLayerNeighborSampler(
        [int(fanout) for fanout in args.fan_out.split(',')])
    dataloader = dgl.sampling.NodeDataLoader(
154
        train_g,
155
156
        train_nid,
        sampler,
157
158
159
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=False,
160
        num_workers=args.num_workers)
161
162

    # Define model and optimizer
163
164
    model = SAGE(in_feats, args.num_hidden, n_classes, args.num_layers, F.relu, args.dropout)
    model = model.to(device)
165
    loss_fcn = nn.CrossEntropyLoss()
166
    loss_fcn = loss_fcn.to(device)
167
168
169
170
171
172
173
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    # Training loop
    avg = 0
    iter_tput = []
    for epoch in range(args.num_epochs):
        tic = time.time()
174
175
176

        # Loop over the dataloader to sample the computation dependency graph as a list of
        # blocks.
177
        for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
178
            tic_step = time.time()
179
180

            # Load the input features as well as output labels
181
            batch_inputs, batch_labels = load_subtensor(train_g, seeds, input_nodes, device)
182
183
184
185
186
187
188
189

            # Compute loss and prediction
            batch_pred = model(blocks, batch_inputs)
            loss = loss_fcn(batch_pred, batch_labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

190
191
            iter_tput.append(len(seeds) / (time.time() - tic_step))
            if step % args.log_every == 0:
192
                acc = compute_acc(batch_pred, batch_labels)
193
                gpu_mem_alloc = th.cuda.max_memory_allocated() / 1000000 if th.cuda.is_available() else 0
194
                print('Epoch {:05d} | Step {:05d} | Loss {:.4f} | Train Acc {:.4f} | Speed (samples/sec) {:.4f} | GPU {:.1f} MiB'.format(
195
                    epoch, step, loss.item(), acc.item(), np.mean(iter_tput[3:]), gpu_mem_alloc))
196
197

        toc = time.time()
198
199
200
201
        print('Epoch Time(s): {:.4f}'.format(toc - tic))
        if epoch >= 5:
            avg += toc - tic
        if epoch % args.eval_every == 0 and epoch != 0:
202
            eval_acc = evaluate(model, val_g, val_g.ndata['features'], val_g.ndata['labels'], val_nid, args.batch_size, device)
203
            print('Eval Acc {:.4f}'.format(eval_acc))
204
205
            test_acc = evaluate(model, test_g, test_g.ndata['features'], test_g.ndata['labels'], test_nid, args.batch_size, device)
            print('Test Acc: {:.4f}'.format(test_acc))
206
207

    print('Avg epoch time: {}'.format(avg / (epoch - 4)))
208
209
210

if __name__ == '__main__':
    argparser = argparse.ArgumentParser("multi-gpu training")
211
212
    argparser.add_argument('--gpu', type=int, default=0,
        help="GPU device ID. Use -1 for CPU training")
213
    argparser.add_argument('--dataset', type=str, default='reddit')
214
215
216
    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)
217
    argparser.add_argument('--fan-out', type=str, default='10,25')
218
219
220
221
    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)
222
223
224
    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.")
225
226
    argparser.add_argument('--inductive', action='store_true',
        help="Inductive learning setting")
227
228
    args = argparser.parse_args()
    
229
230
231
232
    if args.gpu >= 0:
        device = th.device('cuda:%d' % args.gpu)
    else:
        device = th.device('cpu')
233

234
235
236
237
238
239
    if args.dataset == 'reddit':
        g, n_classes = load_reddit()
    elif args.dataset == 'ogb-product':
        g, n_classes = load_ogb('ogbn-products')
    else:
        raise Exception('unknown dataset')
240

241
    in_feats = g.ndata['features'].shape[1]
242
243
244
245
246
247
248
249
250
251
252

    g = dgl.as_heterograph(g)

    if args.inductive:
        train_g, val_g, test_g = inductive_split(g)
    else:
        train_g = val_g = test_g = g

    prepare_mp(train_g)
    prepare_mp(val_g)
    prepare_mp(test_g)
253
    # Pack data
254
    data = in_feats, n_classes, train_g, val_g, test_g
255

256
    run(args, device, data)