Fix benchmark time measurement (#3928)

* Fix benchmark time measurement

* Reduce batch size for bench_rgcn_homogeneous_ns am

bench_rgcn_homogeneous_ns am data sample size is too small for 1024 batch size
Co-authored-by: Rhett Ying <85214957+Rhett-Ying@users.noreply.github.com>
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>

benchmarks/benchmarks/model_speed/bench_gat_ns.py

@@ -90,6 +90,8 @@ def track_time(data):
     lr = 0.003
     dropout = 0.5
     num_workers = 4
+    iter_start = 3
+    iter_count = 10
 
     train_nid = th.nonzero(g.ndata['train_mask'], as_tuple=True)[0]
 
@@ -112,28 +114,9 @@ def track_time(data):
     loss_fcn = loss_fcn.to(device)
     optimizer = optim.Adam(model.parameters(), lr=lr)
 
-    # dry run
-    for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
-        # Load the input features as well as output labels
-        #batch_inputs, batch_labels = load_subtensor(g, seeds, input_nodes, device)
-        blocks = [block.int().to(device) for block in blocks]
-        batch_inputs = blocks[0].srcdata['features']
-        batch_labels = blocks[-1].dstdata['labels']
-
-        # 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 step >= 3:
-            break
-
     # Training loop
     avg = 0
     iter_tput = []
-    t0 = time.time()
     # Loop over the dataloader to sample the computation dependency graph as a list of
     # blocks.
     for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
@@ -149,9 +132,12 @@ def track_time(data):
         loss.backward()
         optimizer.step()
 
-        if step >= 9:  # time 10 loops
+        # start timer at before iter_start
+        if step == iter_start - 1:
+            t0 = time.time()
+        elif step == iter_count + iter_start - 1:  # time iter_count iterations
             break
 
     t1 = time.time()
 
-    return (t1 - t0) / (step + 1)
+    return (t1 - t0) / iter_count

benchmarks/benchmarks/model_speed/bench_pinsage.py

@@ -377,6 +377,8 @@ def track_time(data):
     num_workers = 0
     hidden_dims = 16
     lr = 3e-5
+    iter_start = 3
+    iter_count = 10
 
     g = dataset[0]
     # Sampler
@@ -403,23 +405,8 @@ def track_time(data):
     opt = torch.optim.Adam(model.parameters(), lr=lr)
 
     model.train()
-    for batch_id, (pos_graph, neg_graph, blocks) in enumerate(dataloader):
-        # Copy to GPU
-        for i in range(len(blocks)):
-            blocks[i] = blocks[i].to(device)
-        pos_graph = pos_graph.to(device)
-        neg_graph = neg_graph.to(device)
-
-        loss = model(pos_graph, neg_graph, blocks).mean()
-        opt.zero_grad()
-        loss.backward()
-        opt.step()
-
-        if batch_id >= 3:
-            break
 
     print("start training...")
-    t0 = time.time()
     # For each batch of head-tail-negative triplets...
     for batch_id, (pos_graph, neg_graph, blocks) in enumerate(dataloader):
         # Copy to GPU
@@ -433,9 +420,12 @@ def track_time(data):
         loss.backward()
         opt.step()
 
-        if batch_id >= 10:  # time 10 loops
+        # start timer at before iter_start
+        if batch_id == iter_start - 1:
+            t0 = time.time()
+        elif batch_id == iter_count + iter_start - 1:  # time iter_count iterations
             break
 
     t1 = time.time()
 
-    return (t1 - t0) / (batch_id + 1)
+    return (t1 - t0) / iter_count

benchmarks/benchmarks/model_speed/bench_rgcn_hetero_ns.py

@@ -249,6 +249,8 @@ def track_time(data):
     dropout = 0.5
     use_self_loop = True
     lr = 0.01
+    iter_start = 3
+    iter_count = 10
 
     hg = dataset[0]
     category = dataset.predict_category
@@ -283,31 +285,13 @@ def track_time(data):
         hg, {category: train_idx}, sampler,
         batch_size=batch_size, shuffle=True, num_workers=4)
 
-    # dry run
-    for i, (input_nodes, seeds, blocks) in enumerate(loader):
-        blocks = [blk.to(device) for blk in blocks]
-        seeds = seeds[category]     # we only predict the nodes with type "category"
-        batch_tic = time.time()
-        emb = embed_layer(blocks[0])
-        lbl = labels[seeds].to(device)
-        emb = {k : e.to(device) for k, e in emb.items()}
-        logits = model(emb, blocks)[category]
-        loss = F.cross_entropy(logits, lbl)
-        loss.backward()
-        optimizer.step()
-        sparse_optimizer.step()
-
-        if i >= 3:
-            break
-
     print("start training...")
     model.train()
     embed_layer.train()
     optimizer.zero_grad()
     sparse_optimizer.zero_grad()
 
-    t0 = time.time()
-    for i, (input_nodes, seeds, blocks) in enumerate(loader):
+    for step, (input_nodes, seeds, blocks) in enumerate(loader):
         blocks = [blk.to(device) for blk in blocks]
         seeds = seeds[category]     # we only predict the nodes with type "category"
         batch_tic = time.time()
@@ -320,9 +304,12 @@ def track_time(data):
         optimizer.step()
         sparse_optimizer.step()
 
-        if i >= 9:  # time 10 loops
+        # start timer at before iter_start
+        if step == iter_start - 1:
+            t0 = time.time()
+        elif step == iter_count + iter_start - 1:  # time iter_count iterations
             break
 
     t1 = time.time()
 
-    return (t1 - t0) / (i + 1)
+    return (t1 - t0) / iter_count

benchmarks/benchmarks/model_speed/bench_rgcn_homogeneous_ns.py

@@ -180,9 +180,11 @@ def track_time(data):
     device = utils.get_bench_device()
 
     if data == 'am':
+        batch_size = 64
         n_bases = 40
         l2norm = 5e-4
     elif data == 'ogbn-mag':
+        batch_size = 1024
         n_bases = 2
         l2norm = 0
     else:
@@ -190,13 +192,14 @@ def track_time(data):
 
     fanouts = [25,15]
     n_layers = 2
-    batch_size = 1024
     n_hidden = 64
     dropout = 0.5
     use_self_loop = True
     lr = 0.01
     low_mem = True
     num_workers = 4
+    iter_start = 3
+    iter_count = 10
 
     hg = dataset[0]
     category = dataset.predict_category
@@ -283,32 +286,11 @@ def track_time(data):
     optimizer = th.optim.Adam(all_params, lr=lr, weight_decay=l2norm)
     emb_optimizer = th.optim.SparseAdam(list(embed_layer.node_embeds.parameters()), lr=lr)
 
-    # dry run
-    for i, sample_data in enumerate(loader):
-        input_nodes, output_nodes, blocks = sample_data
-        feats = embed_layer(input_nodes,
-                            blocks[0].srcdata['ntype'],
-                            blocks[0].srcdata['type_id'],
-                            node_feats)
-        logits = model(blocks, feats)
-        seed_idx = blocks[-1].dstdata['type_id']
-        loss = F.cross_entropy(logits, labels[seed_idx])
-        optimizer.zero_grad()
-        emb_optimizer.zero_grad()
-
-        loss.backward()
-        optimizer.step()
-        emb_optimizer.step()
-
-        if i >= 3:
-            break
-
     print("start training...")
     model.train()
     embed_layer.train()
 
-    t0 = time.time()
-    for i, sample_data in enumerate(loader):
+    for step, sample_data in enumerate(loader):
         input_nodes, output_nodes, blocks = sample_data
         feats = embed_layer(input_nodes,
                             blocks[0].srcdata['ntype'],
@@ -324,8 +306,12 @@ def track_time(data):
         optimizer.step()
         emb_optimizer.step()
         
-        if i >= 9:  # time 10 loops
+        # start timer at before iter_start
+        if step == iter_start - 1:
+            t0 = time.time()
+        elif step == iter_count + iter_start - 1:  # time iter_count iterations
             break
+
     t1 = time.time()
 
-    return (t1 - t0) / (i + 1)
+    return (t1 - t0) / iter_count

benchmarks/benchmarks/model_speed/bench_sage_ns.py

@@ -70,6 +70,8 @@ def track_time(data):
     lr = 0.003
     dropout = 0.5
     num_workers = 4
+    iter_start = 3
+    iter_count = 10
 
     train_nid = th.nonzero(g.ndata['train_mask'], as_tuple=True)[0]
 
@@ -92,28 +94,9 @@ def track_time(data):
     loss_fcn = loss_fcn.to(device)
     optimizer = optim.Adam(model.parameters(), lr=lr)
 
-    # dry run
-    for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
-        # Load the input features as well as output labels
-        #batch_inputs, batch_labels = load_subtensor(g, seeds, input_nodes, device)
-        blocks = [block.int().to(device) for block in blocks]
-        batch_inputs = blocks[0].srcdata['features']
-        batch_labels = blocks[-1].dstdata['labels']
-
-        # 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 step >= 3:
-            break
-
     # Training loop
     avg = 0
     iter_tput = []
-    t0 = time.time()
     for step, (input_nodes, seeds, blocks) in enumerate(dataloader):
         # Load the input features as well as output labels
         #batch_inputs, batch_labels = load_subtensor(g, seeds, input_nodes, device)
@@ -128,9 +111,12 @@ def track_time(data):
         loss.backward()
         optimizer.step()
 
-        if step >= 9:  # time 10 loops
+        # start timer at before iter_start
+        if step == iter_start - 1:
+            t0 = time.time()
+        elif step == iter_count + iter_start - 1:  # time iter_count iterations
             break
 
     t1 = time.time()
 
-    return (t1 - t0) / (step + 1)
+    return (t1 - t0) / iter_count

benchmarks/benchmarks/model_speed/bench_sage_unsupervised_ns.py

@@ -114,6 +114,8 @@ def track_time(data, num_negs, batch_size):
     dropout = 0.5
     num_workers = 4
     num_negs = 2
+    iter_start = 3
+    iter_count = 10
 
     n_edges = g.number_of_edges()
     train_seeds = np.arange(n_edges)
@@ -141,28 +143,9 @@ def track_time(data, num_negs, batch_size):
     loss_fcn = loss_fcn.to(device)
     optimizer = optim.Adam(model.parameters(), lr=lr)
 
-    # dry run
-    for step, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(dataloader):
-        # Load the input features as well as output labels
-        batch_inputs = load_subtensor(g, input_nodes, device)
-
-        pos_graph = pos_graph.to(device)
-        neg_graph = neg_graph.to(device)
-        blocks = [block.int().to(device) for block in blocks]
-        # Compute loss and prediction
-        batch_pred = model(blocks, batch_inputs)
-        loss = loss_fcn(batch_pred, pos_graph, neg_graph)
-        optimizer.zero_grad()
-        loss.backward()
-        optimizer.step()
-
-        if step >= 3:
-            break
-
     # Training loop
     avg = 0
     iter_tput = []
-    t0 = time.time()
     for step, (input_nodes, pos_graph, neg_graph, blocks) in enumerate(dataloader):
         # Load the input features as well as output labels
         batch_inputs = load_subtensor(g, input_nodes, device)
@@ -177,9 +160,12 @@ def track_time(data, num_negs, batch_size):
         loss.backward()
         optimizer.step()
 
-        if step >= 9:  # time 10 loops
+        # start timer at before iter_start
+        if step == iter_start - 1:
+            t0 = time.time()
+        elif step == iter_count + iter_start - 1:  # time iter_count iterations
             break
 
     t1 = time.time()
 
-    return (t1 - t0) / (step + 1)
+    return (t1 - t0) / iter_count