more timing

marius-team · Nov 19, 2023 · 0267435 · 0267435
1 parent 417b2a4
commit 0267435
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/src/cpp/src/data/batch.cpp b/src/cpp/src/data/batch.cpp
@@ -145,6 +145,8 @@ void Batch::remoteTo(shared_ptr<c10d::ProcessGroupGloo> pg, int worker_id, int t
 }
 
 void Batch::remoteReceive(shared_ptr<c10d::ProcessGroupGloo> pg, int worker_id, int tag, bool receive_meta) {
+    Timer t_full = new Timer(false);
+    t_full.start();
 
     if (receive_meta) {
         torch::Tensor metadata = torch::tensor({-1, -1, -1}, {torch::kInt32});
@@ -168,6 +170,8 @@ void Batch::remoteReceive(shared_ptr<c10d::ProcessGroupGloo> pg, int worker_id,
         }
         return;
     }
+    Timer t = new Timer(false);
+    t.start();
 
     edges_ = receive_tensor(pg, worker_id, tag);
 
@@ -210,6 +214,11 @@ void Batch::remoteReceive(shared_ptr<c10d::ProcessGroupGloo> pg, int worker_id,
     inv_neg_scores_ = receive_tensor(pg, worker_id, tag);
 
     y_pred_ = receive_tensor(pg, worker_id, tag);
+    t.stop();
+    std::cout<<"batch recv: "<<t.getDuration()<<"\n";
+
+    t_full.stop();
+    std::cout<<"batch recv full: "<<t_full.getDuration()<<"\n";
 }
 
 void Batch::accumulateGradients(float learning_rate) {