Skip to content

B
BTS-MTGNN
Commit 45e9fbb9 by xxx
Options

add time count

parent c8c62f90
Showing with 5 additions and 1 deletions
starrygl/sample/count_static.py
...@@ -32,7 +32,7 @@ class time_count: ...@@ -32,7 +32,7 @@ class time_count:
@staticmethod @staticmethod
def start_gpu(): def start_gpu():
# Uncomment for better breakdown timings # Uncomment for better breakdown timings
torch.cuda.synchronize() #torch.cuda.synchronize()
start_event = torch.cuda.Event(enable_timing=True) start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True)
start_event.record() start_event.record()
......
starrygl/sample/data_loader.py
...@@ -293,6 +293,7 @@ class DistributedDataLoader: ...@@ -293,6 +293,7 @@ class DistributedDataLoader:
t3 = time.time() t3 = time.time()
self.result_queue.append((batch_data,dist_nid,dist_eid,edge_feat,node_feat)) self.result_queue.append((batch_data,dist_nid,dist_eid,edge_feat,node_feat))
tt.time_memory_fetch += tt.elapsed_event(t1) tt.time_memory_fetch += tt.elapsed_event(t1)
#print(tt.time_memory_fetch)
t1 = tt.start_gpu() t1 = tt.start_gpu()
self.submit() self.submit()
tt.time_sample_and_build+=tt.elapsed_event(t1) tt.time_sample_and_build+=tt.elapsed_event(t1)
...@@ -326,6 +327,7 @@ class DistributedDataLoader: ...@@ -326,6 +327,7 @@ class DistributedDataLoader:
node_feat,mem = get_node_feature_by_dist(self.graph,self.mailbox,dist_nid, is_local,out_device=self.device) node_feat,mem = get_node_feature_by_dist(self.graph,self.mailbox,dist_nid, is_local,out_device=self.device)
prepare_input(node_feat,edge_feat,mem,mfgs,dist_nid,dist_eid) prepare_input(node_feat,edge_feat,mem,mfgs,dist_nid,dist_eid)
tt.time_memory_fetch += tt.elapsed_event(t1) tt.time_memory_fetch += tt.elapsed_event(t1)
#print(tt.time_memory_fetch)
if(self.mailbox is not None and self.mailbox.historical_cache is not None): if(self.mailbox is not None and self.mailbox.historical_cache is not None):
id = batch_data[1][0][0].srcdata['ID'] id = batch_data[1][0][0].srcdata['ID']
mask = DistIndex(id).is_shared mask = DistIndex(id).is_shared
...@@ -363,6 +365,8 @@ class DistributedDataLoader: ...@@ -363,6 +365,8 @@ class DistributedDataLoader:
batch_data[1][0][0].srcdata['his_mem'] = batch_data[1][0][0].srcdata['mem'].clone() batch_data[1][0][0].srcdata['his_mem'] = batch_data[1][0][0].srcdata['mem'].clone()
batch_data[1][0][0].srcdata['his_ts'] = batch_data[1][0][0].srcdata['mail_ts'].clone() batch_data[1][0][0].srcdata['his_ts'] = batch_data[1][0][0].srcdata['mail_ts'].clone()
tt.time_memory_fetch+=tt.elapsed_event(t1) tt.time_memory_fetch+=tt.elapsed_event(t1)
#print(tt.time_memory_fetch)
#if(self.mailbox is not None and self.mailbox.historical_cache is not None): #if(self.mailbox is not None and self.mailbox.historical_cache is not None):
# id = batch_data[1][0][0].srcdata['ID'] # id = batch_data[1][0][0].srcdata['ID']
# mask = DistIndex(id).is_shared # mask = DistIndex(id).is_shared
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment