Merge pull request #145 from vrtex-team/hg_sampler

[WIP] Implement HGSampling algorithm from the Heterogeneous Graph Transformer paper

csrc/cpu/hg_sample_cpu.cpp

+#include "hg_sample_cpu.h"
+
+#include <unordered_map>
+#include <random>
+#include "utils.h"
+
+namespace {
+
+void update_budget(
+	int64_t added_node_idx,
+	const node_t &added_node_type,
+	const c10::Dict<int64_t, torch::Tensor> &rowptr_store,
+	const c10::Dict<int64_t, torch::Tensor> &col_store,
+	const c10::Dict<int64_t, edge_t> &edge_type_idx_to_name,
+	std::unordered_map<node_t, std::unordered_set<int64_t>> &sampled_nodes_store,
+	std::unordered_map<node_t, std::unordered_map<int64_t, float>> *budget_store
+) {
+	for (const auto &i : rowptr_store) {
+		const auto &edge_type_idx = i.key();
+		const auto &edge_type = edge_type_idx_to_name.at(edge_type_idx);
+		const auto &source_node_type = std::get<0>(edge_type);
+		const auto &dest_node_type = std::get<2>(edge_type);
+
+		// Skip processing the (rowptr, col) if the node types do not match
+		if (added_node_type.compare(dest_node_type) != 0) {
+			continue;
+		}
+
+		int64_t *row_ptr_raw = i.value().data_ptr<int64_t>();
+		int64_t *col_raw = col_store.at(edge_type_idx).data_ptr<int64_t>();
+
+		// Get the budget map and sampled_nodes for the source node type of the relation
+		const std::unordered_set<int64_t> &sampled_nodes = sampled_nodes_store[source_node_type];
+		std::unordered_map<int64_t, float> &budget = (*budget_store)[source_node_type];
+
+		int64_t row_start_idx = row_ptr_raw[added_node_idx];
+		int64_t row_end_idx = row_ptr_raw[added_node_idx + 1];
+		if (row_start_idx != row_end_idx) {
+			// Compute the norm of degree and update the budget for the neighbors of added_node_idx
+			float norm_deg = 1. / (float)(row_end_idx - row_start_idx);
+			for (int64_t j = row_start_idx; j < row_end_idx; j++) {
+				if (sampled_nodes.find(col_raw[j]) == sampled_nodes.end()) {
+					budget[col_raw[j]] += norm_deg;
+				}
+			}
+		}
+	}
+}
+
+// Sample n nodes according to its type budget map. The probability that node i is sampled is calculated by
+// prob[i] = budget[i]^2 / l2_norm(budget)^2.
+std::unordered_set<int64_t> sample_nodes(const std::unordered_map<int64_t, float> &budget, int n) {	
+	// Compute the squared L2 norm	
+	float norm = 0.0;
+	for (const auto &i : budget) {
+		norm += i.second * i.second;	
+	}
+
+	// Generate n sorted random values between 0 and norm
+	std::vector<float> samples(n);
+	std::uniform_real_distribution<float> dist(0.0, norm);
+	std::default_random_engine gen{std::random_device{}()};
+	std::generate(std::begin(samples), std::end(samples), [&]{ return dist(gen); });
+	std::sort(samples.begin(), samples.end());
+
+	// Iterate through the budget map to compute the cumulative probability cum_prob[i] for node_i. The j-th
+	// sample is assigned to node_i iff cum_prob[i-1] < samples[j] < cum_prob[i]. The implementation assigns
+	// two iterators on budget and samples respectively, then computes the node samples in linear time by 
+	// alternatingly incrementing the two iterators based on their values.
+	std::unordered_set<int64_t> sampled_nodes;
+	sampled_nodes.reserve(samples.size());
+	auto j = samples.begin();
+	float cum_prob = 0.0;
+	for (const auto &i : budget) {
+		cum_prob += i.second * i.second;
+
+		// Increment iterator j until its value is greater than the current cum_prob
+		while (*j < cum_prob && j != samples.end()) {
+			sampled_nodes.insert(i.first);
+			j++;
+		}
+
+		// Terminate early after we complete the sampling
+		if (j == samples.end()) {
+			break;
+		}
+	}
+
+	return sampled_nodes;
+}
+
+}  // namespace
+
+// TODO: Add the appropriate return type
+void hg_sample_cpu(
+	const c10::Dict<int64_t, torch::Tensor> &rowptr_store,
+	const c10::Dict<int64_t, torch::Tensor> &col_store,  
+	const c10::Dict<node_t, torch::Tensor> &origin_nodes_store,
+	const c10::Dict<int64_t, edge_t> &edge_type_idx_to_name,
+	int n,
+	int num_layers
+) {
+	// Verify input
+	for (const auto &kv : rowptr_store) {
+		CHECK_CPU(kv.value());
+	}
+	
+	for (const auto &kv : col_store) {
+		CHECK_CPU(kv.value());
+	}
+	
+	for (const auto &kv : origin_nodes_store) {
+		CHECK_CPU(kv.value());
+	  	CHECK_INPUT(kv.value().dim() == 1);
+	}
+	
+	// Initialize various data structures for the sampling process
+	std::unordered_map<node_t, std::unordered_set<int64_t>> sampled_nodes_store;
+	for (const auto &kv : origin_nodes_store) {
+		const auto &node_type = kv.key();
+		const auto &origin_nodes = kv.value();
+		const int64_t *raw_origin_nodes = origin_nodes.data_ptr<int64_t>();
+
+		// Add each origin node to the sampled_nodes_store
+		for (int64_t i = 0; i < origin_nodes.numel(); i++) {
+			sampled_nodes_store[node_type].insert(raw_origin_nodes[i]);
+		}
+	}
+
+	std::unordered_map<node_t, std::unordered_map<int64_t, float>> budget_store;
+	for (const auto &kv : origin_nodes_store) {
+		const node_t &node_type = kv.key();
+		const auto &origin_nodes = kv.value();
+		const int64_t *raw_origin_nodes = origin_nodes.data_ptr<int64_t>();
+
+		// Update budget for each origin node
+		for (int64_t i = 0; i < origin_nodes.numel(); i++) {
+			update_budget(
+				raw_origin_nodes[i],
+				node_type,
+				rowptr_store,
+				col_store,
+				edge_type_idx_to_name,
+				sampled_nodes_store,
+				&budget_store
+			);
+		}
+	}
+
+
+	// Sampling process
+	for (int l = 0; l < num_layers; l++) {	
+		for (auto &i : budget_store) {
+			const auto &node_type = i.first;
+			auto &budget = i.second;
+			auto &sampled_nodes = sampled_nodes_store[node_type];
+
+			// Perform sampling
+			std::unordered_set<int64_t> new_samples = sample_nodes(budget, n);
+
+			// Remove sampled nodes from the budget and add them to the sampled node store
+			for (const auto &sample : new_samples) {
+				sampled_nodes.insert(sample);
+				budget.erase(sample);
+			}
+
+			// Update the budget
+			for (const auto &sample : new_samples) {
+				update_budget(
+					sample,
+					node_type,
+					rowptr_store,
+					col_store,
+					edge_type_idx_to_name,
+					sampled_nodes_store,
+					&budget_store
+				);
+			}
+		}
+	}
+
+	// Re-index
+	c10::Dict<std::string, std::vector<int64_t>> type_to_n_ids;
+
+}

csrc/cpu/hg_sample_cpu.h

+#pragma once
+
+#include <torch/extension.h>
+
+// Node type is a string and the edge type is a triplet of string representing 
+// (source_node_type, relation_type, dest_node_type).
+typedef std::string node_t;
+typedef std::tuple<std::string, std::string, std::string> edge_t;
+
+// As of PyTorch 1.9.0, c10::Dict does not support tuples or complex data type as key type. We work around this
+// by representing edge types using a single int64_t and a c10::Dict that maps the int64_t index to edge_t.
+void hg_sample_cpu(
+	const c10::Dict<int64_t, torch::Tensor> &rowptr_store,
+	const c10::Dict<int64_t, torch::Tensor> &col_store,  
+	const c10::Dict<node_t, torch::Tensor> &origin_nodes_store,
+	const c10::Dict<int64_t, edge_t> &edge_type_idx_to_name,
+	int n,
+	int num_layers
+);