train_sampling_multi_gpu.py 11.2 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
import dgl.function as fn
import dgl.nn.pytorch as dglnn
import time
import argparse
from dgl.data import RedditDataset
from torch.nn.parallel import DistributedDataParallel
import tqdm
16
import traceback
17

18
from utils import thread_wrapped_func
19
from load_graph import load_reddit, 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
62
63
64
65
66
67
68
69
70
71
72
73
74
        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.

        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?
        nodes = th.arange(g.number_of_nodes())
        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)

75
76
77
78
79
80
81
82
83
84
85
86
            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]
87

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

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

            x = y
        return y

100
101
102
103
104
105
106
107
108
109
110
111
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.

    This is a workaround before full shared memory support on heterogeneous graphs.
    """
    g.in_degree(0)
    g.out_degree(0)
    g.find_edges([0])

112
113
114
115
116
117
def compute_acc(pred, labels):
    """
    Compute the accuracy of prediction given the labels.
    """
    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 : A node ID tensor indicating which nodes do we actually compute the accuracy for.
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, labels, seeds, input_nodes, dev_id):
135
136
137
    """
    Copys features and labels of a set of nodes onto GPU.
    """
138
    batch_inputs = g.ndata['features'][input_nodes].to(dev_id)
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
    batch_labels = labels[seeds].to(dev_id)
    return batch_inputs, batch_labels

#### Entry point

def run(proc_id, n_gpus, args, devices, data):
    # Start up distributed training, if enabled.
    dev_id = devices[proc_id]
    if n_gpus > 1:
        dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
            master_ip='127.0.0.1', master_port='12345')
        world_size = n_gpus
        th.distributed.init_process_group(backend="nccl",
                                          init_method=dist_init_method,
                                          world_size=world_size,
154
                                          rank=proc_id)
155
156
157
    th.cuda.set_device(dev_id)

    # Unpack data
158
159
160
161
162
163
164
    in_feats, n_classes, train_g, val_g, test_g = data
    train_mask = train_g.ndata['train_mask']
    val_mask = val_g.ndata['val_mask']
    test_mask = ~(train_g.ndata['train_mask'] | val_g.ndata['val_mask'])
    train_nid = train_mask.nonzero()[:, 0]
    val_nid = val_mask.nonzero()[:, 0]
    test_nid = test_mask.nonzero()[:, 0]
165
166

    # Split train_nid
167
    train_nid = th.split(train_nid, len(train_nid) // n_gpus)[proc_id]
168

169
    # Create PyTorch DataLoader for constructing blocks
170
171
172
    sampler = dgl.sampling.MultiLayerNeighborSampler(
        [int(fanout) for fanout in args.fan_out.split(',')])
    dataloader = dgl.sampling.NodeDataLoader(
173
        train_g,
174
175
        train_nid,
        sampler,
176
177
178
        batch_size=args.batch_size,
        shuffle=True,
        drop_last=False,
179
        num_workers=args.num_workers)
180
181

    # Define model and optimizer
182
    model = SAGE(in_feats, args.num_hidden, n_classes, args.num_layers, F.relu, args.dropout)
183
184
185
186
187
188
189
190
191
192
193
194
    model = model.to(dev_id)
    if n_gpus > 1:
        model = DistributedDataParallel(model, device_ids=[dev_id], output_device=dev_id)
    loss_fcn = nn.CrossEntropyLoss()
    loss_fcn = loss_fcn.to(dev_id)
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    # Training loop
    avg = 0
    iter_tput = []
    for epoch in range(args.num_epochs):
        tic = time.time()
195
196
197

        # Loop over the dataloader to sample the computation dependency graph as a list of
        # blocks.
198
        for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
199
200
201
202
            if proc_id == 0:
                tic_step = time.time()

            # Load the input features as well as output labels
203
            batch_inputs, batch_labels = load_subtensor(train_g, train_g.ndata['labels'], seeds, input_nodes, dev_id)
204
205
206
207
208
209
210
211
212
213
214
215

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

            if proc_id == 0:
                iter_tput.append(len(seeds) * n_gpus / (time.time() - tic_step))
            if step % args.log_every == 0 and proc_id == 0:
                acc = compute_acc(batch_pred, batch_labels)
216
217
                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:]), th.cuda.max_memory_allocated() / 1000000))
218
219
220
221
222
223
224
225
226
227

        if n_gpus > 1:
            th.distributed.barrier()

        toc = time.time()
        if proc_id == 0:
            print('Epoch Time(s): {:.4f}'.format(toc - tic))
            if epoch >= 5:
                avg += toc - tic
            if epoch % args.eval_every == 0 and epoch != 0:
228
                if n_gpus == 1:
229
230
231
232
                    eval_acc = evaluate(
                        model, val_g, val_g.ndata['features'], val_g.ndata['labels'], val_nid, args.batch_size, devices[0])
                    test_acc = evaluate(
                        model, test_g, test_g.ndata['features'], test_g.ndata['labels'], test_nid, args.batch_size, devices[0])
233
                else:
234
235
236
237
                    eval_acc = evaluate(
                        model.module, val_g, val_g.ndata['features'], val_g.ndata['labels'], val_nid, args.batch_size, devices[0])
                    test_acc = evaluate(
                        model.module, test_g, test_g.ndata['features'], test_g.ndata['labels'], test_nid, args.batch_size, devices[0])
238
                print('Eval Acc {:.4f}'.format(eval_acc))
239
                print('Test Acc: {:.4f}'.format(test_acc))
240

241

242
243
244
245
246
247
248
    if n_gpus > 1:
        th.distributed.barrier()
    if proc_id == 0:
        print('Avg epoch time: {}'.format(avg / (epoch - 4)))

if __name__ == '__main__':
    argparser = argparse.ArgumentParser("multi-gpu training")
249
250
    argparser.add_argument('--gpu', type=str, default='0',
        help="Comma separated list of GPU device IDs.")
251
252
253
    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)
254
    argparser.add_argument('--fan-out', type=str, default='10,25')
255
256
257
258
    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)
259
260
261
    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.")
262
263
    argparser.add_argument('--inductive', action='store_true',
        help="Inductive learning setting")
264
265
266
267
268
    args = argparser.parse_args()
    
    devices = list(map(int, args.gpu.split(',')))
    n_gpus = len(devices)

269
    g, n_classes = load_reddit()
270
    # Construct graph
271
272
273
274
275
276
277
278
279
280
281
    g = dgl.as_heterograph(g)
    in_feats = g.ndata['features'].shape[1]

    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)
282
    # Pack data
283
    data = in_feats, n_classes, train_g, val_g, test_g
284
285
286
287
288
289

    if n_gpus == 1:
        run(0, n_gpus, args, devices, data)
    else:
        procs = []
        for proc_id in range(n_gpus):
290
291
            p = mp.Process(target=thread_wrapped_func(run),
                           args=(proc_id, n_gpus, args, devices, data))
292
293
294
295
            p.start()
            procs.append(p)
        for p in procs:
            p.join()