Merge remote-tracking branch 'dsizov/master'

daniil-sizov · Jul 4, 2023 · 435fc7c · 435fc7c
2 parents 9249be2 + 37de5de
commit 435fc7c
Show file tree

Hide file tree

Showing 9 changed files with 1,204 additions and 146 deletions.
diff --git a/benchmarks/benchmarks/model_speed/bench_gat_ns.py b/benchmarks/benchmarks/model_speed/bench_gat_ns.py
@@ -125,7 +125,7 @@ def track_time(data):
         shuffle=True,
         drop_last=False,
         num_workers=num_workers,
-    )
+        use_cpu_worker_affinity=True)
 
     # Define model and optimizer
     model = GAT(

diff --git a/benchmarks/benchmarks/model_speed/bench_rgcn_hetero_ns.py b/benchmarks/benchmarks/model_speed/bench_rgcn_hetero_ns.py
@@ -345,14 +345,27 @@ def track_time(data):
     )
 
     sampler = dgl.dataloading.MultiLayerNeighborSampler([fanout] * n_layers)
-    loader = dgl.dataloading.DataLoader(
-        hg,
-        {category: train_idx},
-        sampler,
-        batch_size=batch_size,
-        shuffle=True,
-        num_workers=4,
-    )
+    loader = dgl.dataloading.NodeDataLoader(
+        hg, {category: train_idx}, sampler,
+        batch_size=batch_size, shuffle=True, num_workers=4,
+        use_cpu_worker_affinity=True)
+
+    # 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()

diff --git a/benchmarks/benchmarks/model_speed/bench_rgcn_homogeneous_ns.py b/benchmarks/benchmarks/model_speed/bench_rgcn_homogeneous_ns.py
@@ -290,7 +290,7 @@ def track_time(data):
         shuffle=True,
         drop_last=False,
         num_workers=num_workers,
-    )
+        use_cpu_worker_affinity=True)
 
     # node features
     # None for one-hot feature, if not none, it should be the feature tensor.

diff --git a/benchmarks/benchmarks/model_speed/bench_sage_ns.py b/benchmarks/benchmarks/model_speed/bench_sage_ns.py
@@ -87,7 +87,7 @@ def track_time(data):
         shuffle=True,
         drop_last=False,
         num_workers=num_workers,
-    )
+        use_cpu_worker_affinity=True)
 
     # Define model and optimizer
     model = SAGE(in_feats, num_hidden, n_classes, num_layers, F.relu, dropout)