fix time count

examples/train_boundery.py

@@ -562,6 +562,14 @@ def main():
             optimizer.zero_grad()
             ones = torch.ones(metadata['dst_neg_index'].shape[0],device = model.device,dtype=torch.float)
             pred_pos, pred_neg = model(mfgs,metadata,neg_samples=args.neg_samples,async_param = param)
+            ada_param.update_gnn_aggregate_time(ada_param.last_start_event_gnn_aggregate)
+            if len(trainloader.result_queue) > 0:
+                batch_data,dist_nid,dist_eid,edge_feat,node_feat0 = trainloader.result_queue[0]
+                edge_feat[1].wait()
+                node_feat0[1].wait()
+                if ada_param is not None:
+                    ada_param.update_fetch_time(ada_param.last_start_event_fetch)
+                    ada_param.update_parameter()
             #print(time_count.elapsed_event(t2))
             loss = creterion(pred_pos, torch.ones_like(pred_pos)) 
             if args.local_neg_sample is False:
@@ -588,7 +596,7 @@ def main():
             mailbox.update_p2p_mem()
             mailbox.update_p2p_mail()
             start = time_count.start_gpu()
-            ada_param.update_gnn_aggregate_time(ada_param.last_start_event_gnn_aggregate)
+            
             #ada_param.update_parameter()
             #torch.cuda.empty_cache()
 

starrygl/module/utils.py

@@ -141,9 +141,9 @@ class AdaParameter:
     if self.end_event_fetch is None or self.end_event_memory_sync is None or self.end_event_memory_update is None or self.end_event_gnn_aggregate is None:
       return
     self.end_event_fetch[1].synchronize()
-    self.end_event_memory_sync[1].synchronize()
-    self.end_event_memory_update[1].synchronize()
-    self.end_event_gnn_aggregate[1].synchronize()
+    #self.end_event_memory_sync[1].synchronize()
+    #self.end_event_memory_update[1].synchronize()
+    #self.end_event_gnn_aggregate[1].synchronize()
     self.average_fetch += self.end_event_fetch[0].elapsed_time(self.end_event_fetch[1])
     self.count_fetch += 1
     self.average_memory_sync += self.end_event_memory_sync[0].elapsed_time(self.end_event_memory_sync[1])

starrygl/sample/data_loader.py

@@ -392,9 +392,7 @@ class DistributedDataLoader:
                             node_feat0 = node_feat0[0]
                             node_feat = None
                             mem = self.mailbox.unpack(node_feat0,mailbox = True)
-                        if self.ada_param is not None:
-                            self.ada_param.update_fetch_time(self.ada_param.last_start_event_fetch)
-                            self.ada_param.update_parameter()
+                        
                         #print(node_feat.shape,edge_feat.shape,mem[0].shape)
                         #node_feat[1].wait()
                         #node_feat = node_feat[0]