train_sampling_multi_gpu.py 10.5 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
            h = layer(block, h)
45
46
47
            if l != len(self.layers) - 1:
                h = self.activation(h)
                h = self.dropout(h)
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
        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)

68
69
70
71
72
73
74
75
76
77
78
79
            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]
80

81
                block = block.to(device)
82
                h = x[input_nodes].to(device)
83
                h = layer(block, h)
84
85
86
                if l != len(self.layers) - 1:
                    h = self.activation(h)
                    h = self.dropout(h)
87

88
                y[output_nodes] = h.cpu()
89
90
91
92
93
94
95
96
97
98

            x = y
        return y

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

99
def evaluate(model, g, inputs, labels, val_nid, batch_size, device):
100
    """
101
    Evaluate the model on the validation set specified by ``val_nid``.
102
103
104
    g : The entire graph.
    inputs : The features of all the nodes.
    labels : The labels of all the nodes.
105
    val_nid : A node ID tensor indicating which nodes do we actually compute the accuracy for.
106
107
108
109
110
111
112
    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()
113
    return compute_acc(pred[val_nid], labels[val_nid])
114

115
def load_subtensor(g, labels, seeds, input_nodes, dev_id):
116
117
118
    """
    Copys features and labels of a set of nodes onto GPU.
    """
119
    batch_inputs = g.ndata['features'][input_nodes].to(dev_id)
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    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,
135
                                          rank=proc_id)
136
137
138
    th.cuda.set_device(dev_id)

    # Unpack data
139
140
141
    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']
142
    test_mask = ~(test_g.ndata['train_mask'] | test_g.ndata['val_mask'])
143
144
145
    train_nid = train_mask.nonzero()[:, 0]
    val_nid = val_mask.nonzero()[:, 0]
    test_nid = test_mask.nonzero()[:, 0]
146
147

    # Split train_nid
148
    train_nid = th.split(train_nid, len(train_nid) // n_gpus)[proc_id]
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
    model = SAGE(in_feats, args.num_hidden, n_classes, args.num_layers, F.relu, args.dropout)
164
165
166
167
168
169
170
171
172
173
174
175
    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()
176
177
178

        # Loop over the dataloader to sample the computation dependency graph as a list of
        # blocks.
179
        for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
180
181
182
183
            if proc_id == 0:
                tic_step = time.time()

            # Load the input features as well as output labels
184
            batch_inputs, batch_labels = load_subtensor(train_g, train_g.ndata['labels'], seeds, input_nodes, dev_id)
185
            blocks = [block.to(dev_id) for block in blocks]
186
187
188
189
190
191
192
193
194
195
196
            # 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)
197
198
                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))
199
200
201
202
203
204
205
206
207
208

        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:
209
                if n_gpus == 1:
210
211
212
213
                    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])
214
                else:
215
216
217
218
                    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])
219
                print('Eval Acc {:.4f}'.format(eval_acc))
220
                print('Test Acc: {:.4f}'.format(test_acc))
221

222

223
224
225
226
227
228
229
    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")
230
231
    argparser.add_argument('--gpu', type=str, default='0',
        help="Comma separated list of GPU device IDs.")
232
233
234
    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)
235
    argparser.add_argument('--fan-out', type=str, default='10,25')
236
237
238
239
    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)
240
241
242
    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.")
243
244
    argparser.add_argument('--inductive', action='store_true',
        help="Inductive learning setting")
245
246
247
248
249
    args = argparser.parse_args()
    
    devices = list(map(int, args.gpu.split(',')))
    n_gpus = len(devices)

250
    g, n_classes = load_reddit()
251
    # Construct graph
252
253
254
255
256
257
258
259
    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

260
261
262
    train_g.create_format_()
    val_g.create_format_()
    test_g.create_format_()
263
    # Pack data
264
    data = in_feats, n_classes, train_g, val_g, test_g
265
266
267
268
269
270

    if n_gpus == 1:
        run(0, n_gpus, args, devices, data)
    else:
        procs = []
        for proc_id in range(n_gpus):
271
272
            p = mp.Process(target=thread_wrapped_func(run),
                           args=(proc_id, n_gpus, args, devices, data))
273
274
275
276
            p.start()
            procs.append(p)
        for p in procs:
            p.join()