read feature change

marius-team · Nov 21, 2023 · 1a071e6 · 1a071e6
1 parent 2eecd7d
commit 1a071e6
Showing 1 changed file with 17 additions and 7 deletions.
diff --git a/src/cpp/src/pipeline/pipeline_gpu.cpp b/src/cpp/src/pipeline/pipeline_gpu.cpp
@@ -73,13 +73,13 @@ void RemoteLoadWorker::run() {
             }
 
             batch->remoteReceive(pipeline_->model_->pg_gloo_->pg, parent, tag);
-            if (batch->sub_batches_.size() > 0) {
-                for (int i = 0; i < batch->sub_batches_.size(); i++) {
-                    batch->sub_batches_[i]->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->sub_batches_[i]->unique_node_indices_);
-                }
-            } else {
-                batch->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->unique_node_indices_);
-            }
+//            if (batch->sub_batches_.size() > 0) {
+//                for (int i = 0; i < batch->sub_batches_.size(); i++) {
+//                    batch->sub_batches_[i]->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->sub_batches_[i]->unique_node_indices_);
+//                }
+//            } else {
+//                batch->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->unique_node_indices_);
+//            }
 
             if (pipeline_->model_->device_.is_cuda()) {
                 ((PipelineGPU *)pipeline_)->loaded_batches_->blocking_push(batch);
@@ -150,6 +150,16 @@ void BatchToDeviceWorker::run() {
                 break;
             }
 
+            if (batch->sub_batches_.size() > 0) {
+                for (int i = 0; i < batch->sub_batches_.size(); i++) {
+                    if (!batch->sub_batches_[i]->node_features_.defined())
+                        batch->sub_batches_[i]->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->sub_batches_[i]->unique_node_indices_);
+                }
+            } else {
+                if (!batch->node_features_.defined())
+                    batch->node_features_ = pipeline_->dataloader_->graph_storage_->getNodeFeatures(batch->unique_node_indices_);
+            }
+
             batchToDevice(pipeline_, batch);
             t.stop();
             std::cout<<"batch to: "<<t.getDuration()<<"\n";