Added optimization of biased sampling

pyg-team · Oct 31, 2023 · 225785d · 225785d
1 parent 9fc1afc
commit 225785d
Showing 1 changed file with 17 additions and 5 deletions.
diff --git a/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp b/pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp
@@ -223,11 +223,23 @@ class NeighborSampler {
 
     // Case 2: Multinomial sampling:
     else {
-      const auto index = at::multinomial(weight, count, replace);
-      const auto index_data = index.data_ptr<int64_t>();
-      for (size_t i = 0; i < index.numel(); ++i) {
-        add(row_start + index_data[i], global_src_node, local_src_node,
-            dst_mapper, out_global_dst_nodes);
+      if (replace) {
+        const auto index = at::multinomial(weight, count, replace);
+        const auto index_data = index.data_ptr<int64_t>();
+        for (size_t i = 0; i < index.numel(); ++i) {
+          add(row_start + index_data[i], global_src_node, local_src_node,
+              dst_mapper, out_global_dst_nodes);
+        }
+      }
+      else {
+        const auto rand = at::empty_like(weight).uniform_();
+        const auto a = (rand.log() / weight);
+        const auto index = std::get<1>(a.topk(count));
+        const auto index_data = index.data_ptr<int64_t>();
+        for (size_t i = 0; i < index.numel(); ++i) {
+          add(row_start + index_data[i], global_src_node, local_src_node,
+              dst_mapper, out_global_dst_nodes);
+        }
       }
     }
   }