merge master

.gitattributes

+install.sh merge=ours
\ No newline at end of file

.gitignore

+*.tgz
+*.my
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

csrc/sampler/export.cpp

 #include<head.h>
 #include <sampler.h>
+#include <tppr.h>
 #include <output.h>
 #include <neighbors.h>
 #include <temporal_utils.h>
@@ -88,4 +89,22 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
         .def("reset", &ParallelSampler::reset)
         .def("get_ret", [](const ParallelSampler &ps) { return ps.ret; });
 
+    py::class_<ParallelTppRComputer>(m, "ParallelTppRComputer")
+        .def(py::init<TemporalNeighborBlock &, NodeIDType, EdgeIDType, int,
+                      int, int, int, vector<float>&, vector<float>& >())
+        .def_readonly("ret", &ParallelTppRComputer::ret, py::return_value_policy::reference)
+        .def("reset_ret", &ParallelTppRComputer::reset_ret)
+        .def("reset_tppr", &ParallelTppRComputer::reset_tppr)
+        .def("reset_val_tppr", &ParallelTppRComputer::reset_val_tppr)
+        .def("backup_tppr", &ParallelTppRComputer::backup_tppr)
+        .def("restore_tppr", &ParallelTppRComputer::restore_tppr)
+        .def("restore_val_tppr", &ParallelTppRComputer::restore_val_tppr)
+        .def("get_pruned_topk", &ParallelTppRComputer::get_pruned_topk)
+        .def("extract_streaming_tppr", &ParallelTppRComputer::extract_streaming_tppr)
+        .def("streaming_topk", &ParallelTppRComputer::streaming_topk)
+        .def("single_streaming_topk", &ParallelTppRComputer::single_streaming_topk)
+        .def("streaming_topk_no_fake", &ParallelTppRComputer::streaming_topk_no_fake)
+        .def("compute_val_tppr", &ParallelTppRComputer::compute_val_tppr)
+        .def("get_ret", [](const ParallelTppRComputer &ps) { return ps.ret; });
+
 }
\ No newline at end of file

csrc/sampler/include/head.h

 #pragma once
 #include <iostream>
+#include <algorithm>
 #include <torch/extension.h>
 #include <omp.h>
 #include <time.h>
@@ -17,6 +18,12 @@ typedef int64_t NodeIDType;
 typedef int64_t EdgeIDType;
 typedef float WeightType;
 typedef float TimeStampType;
+typedef tuple<NodeIDType, EdgeIDType, TimeStampType> PPRKeyType;
+typedef double PPRValueType;
+typedef phmap::parallel_flat_hash_map<PPRKeyType, PPRValueType> PPRDictType;
+typedef vector<PPRDictType> PPRListDictType;
+typedef vector<vector<PPRDictType>> PPRListListDictType;
+typedef vector<vector<double>> NormListType;
 
 class TemporalNeighborBlock;
 class TemporalGraphBlock;
@@ -28,6 +35,7 @@ int nodeIdToInOut(NodeIDType nid, int pid, const vector<NodeIDType>& part_ptr);
 int nodeIdToPartId(NodeIDType nid, const vector<NodeIDType>& part_ptr);
 vector<th::Tensor> divide_nodes_to_part(th::Tensor nodes, const vector<NodeIDType>& part_ptr, int threads);
 NodeIDType sample_multinomial(const vector<WeightType>& weights, default_random_engine& e);
+vector<int64_t> sample_max(const vector<WeightType>& weights, int k);
 
 
 
@@ -173,3 +181,17 @@ NodeIDType sample_multinomial(const vector<WeightType>& weights, default_random_
     sample_indice = distance(cumulative_weights.begin(), it);
     return sample_indice;
 }
+
+vector<int64_t> sample_max(const vector<WeightType>& weights, int k) {
+    vector<int64_t> indices(weights.size());
+    for (int i = 0; i < weights.size(); ++i) {
+        indices[i] = i;
+    }
+
+    // 使用部分排序算法（选择算法）找到前k个最大值的索引
+    partial_sort(indices.begin(), indices.begin() + k, indices.end(), 
+                 [&weights](int64_t a, int64_t b) { return weights[a] > weights[b]; });
+
+    // 返回前k个最大值的索引
+    return vector<int64_t>(indices.begin(), indices.begin() + k);
+}
\ No newline at end of file

csrc/sampler/include/output.h

@@ -11,6 +11,7 @@ class TemporalGraphBlock
         vector<int64_t> src_index;
         vector<NodeIDType> sample_nodes;
         vector<TimeStampType> sample_nodes_ts;
+        vector<WeightType> e_weights;
         double sample_time = 0;
         double tot_time = 0;
         int64_t sample_edge_num = 0;

csrc/sampler/include/tppr.h

+#pragma once
+#include <head.h>
+#include <neighbors.h>
+# include <output.h>
+
+class ParallelTppRComputer
+{
+    public:
+        TemporalNeighborBlock& tnb;
+        NodeIDType num_nodes;
+        EdgeIDType num_edges;
+        int threads;
+        int fanout;//k, width
+        int num_layers;//depth
+        int num_tpprs;//n_tpprs
+        vector<float> alpha_list;
+        vector<float> beta_list;
+        // string policy;
+        PPRListListDictType PPR_list;
+        PPRListListDictType val_PPR_list;
+        NormListType norm_list;
+        NormListType val_norm_list;
+        vector<vector<TemporalGraphBlock>> ret;
+
+        ParallelTppRComputer(TemporalNeighborBlock& _tnb, NodeIDType _num_nodes, EdgeIDType _num_edges, int _threads, 
+                        int _fanout, int _num_layers, int _num_tpprs, vector<float>& _alpha_list, vector<float>& _beta_list) :
+                        tnb(_tnb), num_nodes(_num_nodes), num_edges(_num_edges), threads(_threads), 
+                        fanout(_fanout), num_layers(_num_layers), num_tpprs(_num_tpprs), alpha_list(_alpha_list), beta_list(_beta_list)
+        {
+            omp_set_num_threads(_threads);
+            ret.clear();
+            ret = vector<vector<TemporalGraphBlock>>(_num_tpprs, vector<TemporalGraphBlock>());
+        }
+
+        void reset_ret() {
+            for (int i = 0; i < num_tpprs; ++i) {
+                ret[i].clear(); // 清空每个内部的 vector
+            }
+        }
+
+        void reset_ret_i(int tppr_id) {
+            ret[tppr_id].clear(); // 清空 tppr_id 处的 vector
+        }
+        
+        void reset_tppr(){
+            PPR_list = PPRListListDictType(num_tpprs, PPRListDictType(num_nodes));
+            norm_list = NormListType(num_tpprs, vector<double>(num_nodes, 0.0));
+        }
+
+        void reset_val_tppr(){
+            val_PPR_list = PPRListListDictType(num_tpprs, PPRListDictType(num_nodes));
+            val_norm_list = NormListType(num_tpprs, vector<double>(num_nodes, 0.0));
+        }
+        
+        py::tuple backup_tppr(){
+            return py::make_tuple(this->PPR_list, this->norm_list);
+        }
+
+        void restore_tppr(PPRListListDictType& input_PPR_list, NormListType& input_norm_list){
+            this->PPR_list = input_PPR_list;
+            this->norm_list = input_norm_list;
+        }
+
+        void restore_val_tppr(PPRListListDictType& input_PPR_list, NormListType& input_norm_list){
+            this->val_PPR_list = input_PPR_list;
+            this->val_norm_list = input_norm_list;
+        }
+
+        PPRDictType compute_s1_s2(NodeIDType s1, NodeIDType s2, int tppr_id, EdgeIDType eid, TimeStampType ts);
+        void get_pruned_topk(th::Tensor src_nodes, th::Tensor root_ts, int tppr_id);
+        void extract_streaming_tppr(PPRDictType tppr_dict, TimeStampType current_ts, int index0, int position);
+        void streaming_topk(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids);
+        void single_streaming_topk(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids, int tppr_id);
+        void streaming_topk_no_fake(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids);
+        void compute_val_tppr(th::Tensor src_nodes, th::Tensor dst_nodes, th::Tensor root_ts, th::Tensor eids);
+};
+
+PPRDictType ParallelTppRComputer :: compute_s1_s2(NodeIDType s1, NodeIDType s2, int tppr_id, EdgeIDType eid, TimeStampType ts){
+    int alpha = alpha_list[tppr_id], beta = beta_list[tppr_id];
+    vector<double> norm_list = this->norm_list[tppr_id];
+    PPRListDictType PPR_list = this->PPR_list[tppr_id];
+    PPRDictType t_s1_PPR= PPRDictType();
+    PPRDictType updated_tppr= PPRDictType();
+    float scala_s1, scala_s2;
+    /***************s1 side*******************/
+    if(norm_list[s1]==0){
+        scala_s2 = 1-alpha;
+    }
+    else{
+        t_s1_PPR = PPR_list[s1];
+        double last_norm = norm_list[s1], new_norm;
+        new_norm = last_norm*beta+beta;
+        scala_s1 = last_norm/new_norm*beta;
+        scala_s2 = beta/new_norm*(1-alpha);
+        for (const auto& pair : t_s1_PPR)
+            t_s1_PPR[pair.first] = pair.second*scala_s1;
+    }
+    /**************s2 side*******************/
+    if(norm_list[s1]==0){
+        t_s1_PPR[make_tuple(eid, s2, ts)] = alpha!=0 ? scala_s2*alpha : scala_s2;
+    }
+    else{
+        PPRDictType s2_PPR= PPR_list[s2];
+        for (const auto& pair : s2_PPR){
+            if(t_s1_PPR.count(pair.first)==1)
+                t_s1_PPR[pair.first] += pair.second*scala_s2;
+            else
+                t_s1_PPR[pair.first] = pair.second*scala_s2;
+        }
+        t_s1_PPR[make_tuple(eid, s2, ts)] = alpha!=0 ? scala_s2*alpha : scala_s2;
+    }
+    /*********exract the top-k items ********/
+    int tppr_size = t_s1_PPR.size();
+    if(tppr_size<=this->fanout)
+        updated_tppr = t_s1_PPR;
+    else{
+        std::vector<std::pair<PPRKeyType, PPRValueType>> pairs;
+        pairs.reserve(t_s1_PPR.size());
+        // 提取键值对到 pairs 向量
+        for (const auto& pair : t_s1_PPR) {
+            pairs.emplace_back(pair.first, pair.second);
+        }
+        // 使用并行部分排序来获得前 this->fanout 个元素
+        std::partial_sort(pairs.begin(), pairs.begin() + this->fanout, pairs.end(),
+                        [](const auto& a, const auto& b) { return a.second > b.second; });
+        // 将部分排序后的键值对添加到 updated_tppr
+        for (size_t i = 0; i < this->fanout; ++i) {
+            const auto& pair = pairs[i];
+            updated_tppr[pair.first] = pair.second;
+        }
+    }
+    return updated_tppr;
+}
+
+void ParallelTppRComputer :: get_pruned_topk(th::Tensor src_nodes, th::Tensor root_ts, int tppr_id){
+    auto src_nodes_data = get_data_ptr<NodeIDType>(src_nodes);
+    auto ts_data = get_data_ptr<TimeStampType>(root_ts);
+    int64_t n_edges = src_nodes.size(0);
+    float alpha = alpha_list[tppr_id], beta = beta_list[tppr_id];
+    this->reset_ret_i(tppr_id);
+    for(int i=0;i<n_edges;i++)
+    {
+        NodeIDType target_node =  src_nodes_data[i];
+        TimeStampType target_timestamp =  ts_data[i];
+        PPRDictType tppr_dict;
+        /*******get dictionary of neighbors*********************/
+        vector<tuple<NodeIDType, TimeStampType, PPRValueType>> query_list;
+        query_list.push_back(make_tuple(target_node,  target_timestamp, 1.0));
+        for(int depth=0;depth<this->num_layers;depth++)
+        {
+            vector<tuple<NodeIDType, TimeStampType, PPRValueType>> new_query_list;
+            /*******traverse the query list*********************/
+            for(int j=0;j<query_list.size();j++)
+            {
+                NodeIDType query_node = get<0>(query_list[j]);
+                NodeIDType query_ts = get<1>(query_list[j]);
+                NodeIDType query_weight = get<2>(query_list[j]);
+
+                int end_index = lower_bound(tnb.timestamp[query_node].begin(), tnb.timestamp[query_node].end(), query_ts)-tnb.timestamp[query_node].begin();
+                int n_ngh = end_index;
+                if(n_ngh==0) continue;
+                else
+                {
+                    double norm = beta/(1-beta)*(1-pow(beta, n_ngh));
+                    double weight = alpha!=0 && depth==0 ? query_weight*(1-alpha)*beta/norm*alpha : query_weight*(1-alpha)*beta/norm;
+
+                    for(int z=0;z<min(this->fanout, n_ngh);z++){
+                        EdgeIDType eid = tnb.eid[query_node][end_index-z-1];
+                        NodeIDType node = tnb.neighbors[query_node][end_index-z-1];
+                        // the timestamp here is a neighbor timestamp, 
+                        // so that it is indeed a temporal random walk
+                        TimeStampType timestamp =  tnb.timestamp[query_node][end_index-z-1];
+                        PPRKeyType state = make_tuple(eid, node, timestamp);
+                        // update dict
+                        if(tppr_dict.count(state)==1)
+                            tppr_dict[state] = tppr_dict[state]+weight;
+                        else
+                            tppr_dict[state] = weight;
+                        // update query list
+                        tuple<NodeIDType, TimeStampType, PPRValueType> new_query = make_tuple(node, timestamp, weight);
+                        new_query_list.push_back(new_query);
+                        // update weight
+                        weight = weight*beta;
+                    }
+                }
+            }
+            if(new_query_list.empty()) break;
+            else query_list = new_query_list;
+        }
+        /*****sort and get the top-k neighbors********/
+        int tppr_size = tppr_dict.size();
+        if(tppr_size==0) continue;
+
+        TimeStampType current_timestamp = ts_data[i];
+
+        PPRDictType updated_tppr= PPRDictType();            
+        if(tppr_size<=this->fanout)
+            updated_tppr = tppr_dict;
+        else
+        {
+            std::vector<std::pair<PPRKeyType, PPRValueType>> pairs;
+            pairs.reserve(tppr_dict.size());
+            // 提取键值对到 pairs 向量
+            for (const auto& pair : tppr_dict) {
+                pairs.emplace_back(pair.first, pair.second);
+            }
+            // 使用并行部分排序来获得前 this->fanout 个元素
+            std::partial_sort(pairs.begin(), pairs.begin() + this->fanout, pairs.end(),
+                            [](const auto& a, const auto& b) { return a.second > b.second; });
+            // 将部分排序后的键值对添加到 updated_tppr
+            for (size_t i = 0; i < this->fanout; ++i) {
+                const auto& pair = pairs[i];
+                updated_tppr[pair.first] = pair.second;
+            }
+        }
+        // this->PPR_list[tppr_id][target_node] = updated_tppr;
+        extract_streaming_tppr(updated_tppr, current_timestamp, tppr_id, i);
+    }
+}
+
+// category=0-src  category=1-dst category=2-fake
+void ParallelTppRComputer :: extract_streaming_tppr(PPRDictType tppr_dict, TimeStampType current_ts, int index0, int position){
+    ret[index0][position] = TemporalGraphBlock();
+    if(!tppr_dict.empty()){
+        ret[index0][position].sample_nodes.resize(this->fanout);
+        ret[index0][position].eid.resize(this->fanout);
+        ret[index0][position].sample_nodes_ts.resize(this->fanout);
+        ret[index0][position].e_weights.resize(this->fanout);
+        ret[index0][position].delta_ts.resize(this->fanout);
+        int j=0;
+        for (const auto& pair : tppr_dict){
+            auto tuple = pair.first;
+            auto weight = pair.second;
+            EdgeIDType eid = get<0>(tuple);
+            NodeIDType dst = get<1>(tuple);
+            TimeStampType ets = get<2>(tuple);
+
+            ret[index0][position].sample_nodes[j]=dst;
+            ret[index0][position].eid[j]=eid;
+            ret[index0][position].sample_nodes_ts[j]=ets;
+            ret[index0][position].e_weights[j]=weight;
+            ret[index0][position].delta_ts[j]=current_ts-ets;
+
+            j++;
+        }
+    }
+}
+
+void ParallelTppRComputer :: streaming_topk(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids){
+    auto src_nodes_data = get_data_ptr<NodeIDType>(src_nodes);
+    auto ts_data = get_data_ptr<TimeStampType>(root_ts);
+    auto eids_data = get_data_ptr<EdgeIDType>(eids);
+    int n_nodes = src_nodes.size(0);
+    int n_edges = num_nodes/3;
+    this->reset_ret();
+    for(int index0=0;index0<num_tpprs;index0++){
+        int alpha = alpha_list[index0], beta = beta_list[index0];
+        ret[index0].resize(n_nodes);
+        vector<double>& norm_list = this->norm_list[index0];
+        PPRListDictType& PPR_list = this->PPR_list[index0];
+        for(int i=0; i<n_edges; i++){
+            NodeIDType src = src_nodes_data[i];
+            NodeIDType dst = src_nodes_data[i+n_edges];
+            NodeIDType fake = src_nodes_data[i+(n_edges<<1)];
+            TimeStampType ts = ts_data[i];
+            EdgeIDType eid = eids_data[i];
+
+            /******first extract the top-k neighbors and fill the list******/
+            extract_streaming_tppr(PPR_list[src], ts, index0, i);
+            extract_streaming_tppr(PPR_list[dst], ts, index0, i+n_edges);
+            extract_streaming_tppr(PPR_list[fake], ts, index0, i+(n_edges<<1));
+
+            /******then update the PPR values here**************************/
+            PPR_list[src] = compute_s1_s2(src, dst, index0, eid, ts);
+            norm_list[src] = norm_list[src]*beta+beta;
+            if(src!=dst){
+                PPR_list[dst] = compute_s1_s2(dst, src, index0, eid, ts);
+                norm_list[dst] = norm_list[dst]*beta+beta;
+            }
+        }
+    }
+}
+
+void ParallelTppRComputer :: single_streaming_topk(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids, int tppr_id){
+    auto src_nodes_data = get_data_ptr<NodeIDType>(src_nodes);
+    auto ts_data = get_data_ptr<TimeStampType>(root_ts);
+    auto eids_data = get_data_ptr<EdgeIDType>(eids);
+    int n_nodes = src_nodes.size(0);
+    int n_edges = num_nodes/3;
+    this->reset_ret_i(tppr_id);
+    int alpha = alpha_list[tppr_id], beta = beta_list[tppr_id];
+    ret[tppr_id].resize(n_nodes);
+    vector<double>& norm_list = this->norm_list[tppr_id];
+    PPRListDictType& PPR_list = this->PPR_list[tppr_id];
+    for(int i=0; i<n_edges; i++){
+        NodeIDType src = src_nodes_data[i];
+        NodeIDType dst = src_nodes_data[i+n_edges];
+        NodeIDType fake = src_nodes_data[i+(n_edges<<1)];
+        TimeStampType ts = ts_data[i];
+        EdgeIDType eid = eids_data[i];
+
+        /******first extract the top-k neighbors and fill the list******/
+        extract_streaming_tppr(PPR_list[src], ts, tppr_id, i);
+        extract_streaming_tppr(PPR_list[dst], ts, tppr_id, i+n_edges);
+        extract_streaming_tppr(PPR_list[fake], ts, tppr_id, i+(n_edges<<1));
+
+        /******then update the PPR values here**************************/
+        PPR_list[src] = compute_s1_s2(src, dst, tppr_id, eid, ts);
+        norm_list[src] = norm_list[src]*beta+beta;
+        if(src!=dst){
+            PPR_list[dst] = compute_s1_s2(dst, src, tppr_id, eid, ts);
+            norm_list[dst] = norm_list[dst]*beta+beta;
+        }
+    }
+}
+
+void ParallelTppRComputer :: streaming_topk_no_fake(th::Tensor src_nodes, th::Tensor root_ts, th::Tensor eids){
+    auto src_nodes_data = get_data_ptr<NodeIDType>(src_nodes);
+    auto ts_data = get_data_ptr<TimeStampType>(root_ts);
+    auto eids_data = get_data_ptr<EdgeIDType>(eids);
+    int n_nodes = src_nodes.size(0);
+    int n_edges = num_nodes/2;
+    this->reset_ret();
+    for(int index0=0;index0<num_tpprs;index0++){
+        int alpha = alpha_list[index0], beta = beta_list[index0];
+        ret[index0].resize(n_nodes);
+        vector<double>& norm_list = this->norm_list[index0];
+        PPRListDictType& PPR_list = this->PPR_list[index0];
+        for(int i=0; i<n_edges; i++){
+            NodeIDType src = src_nodes_data[i];
+            NodeIDType dst = src_nodes_data[i+n_edges];
+            TimeStampType ts = ts_data[i];
+            EdgeIDType eid = eids_data[i];
+
+            /******first extract the top-k neighbors and fill the list******/
+            extract_streaming_tppr(PPR_list[src], ts, index0, i);
+            extract_streaming_tppr(PPR_list[dst], ts, index0, i+n_edges);
+
+            /******then update the PPR values here**************************/
+            PPR_list[src] = compute_s1_s2(src, dst, index0, eid, ts);
+            norm_list[src] = norm_list[src]*beta+beta;
+            if(src!=dst){
+                PPR_list[dst] = compute_s1_s2(dst, src, index0, eid, ts);
+                norm_list[dst] = norm_list[dst]*beta+beta;
+            }
+        }
+    }
+}
+
+void ParallelTppRComputer :: compute_val_tppr(th::Tensor src_nodes, th::Tensor dst_nodes, th::Tensor root_ts, th::Tensor eids){
+    auto src_nodes_data = get_data_ptr<NodeIDType>(src_nodes);
+    auto dst_nodes_data = get_data_ptr<NodeIDType>(dst_nodes);
+    auto ts_data = get_data_ptr<TimeStampType>(root_ts);
+    auto eids_data = get_data_ptr<EdgeIDType>(eids);
+    int n_edges = src_nodes.size(0);
+    for(int index0=0;index0<num_tpprs;index0++){
+        int alpha = alpha_list[index0], beta = beta_list[index0];
+        vector<double>& norm_list = this->norm_list[index0];
+        PPRListDictType& PPR_list = this->PPR_list[index0];
+        for(int i=0; i<n_edges; i++){
+            NodeIDType src = src_nodes_data[i];
+            NodeIDType dst = dst_nodes_data[i];
+            TimeStampType ts = ts_data[i];
+            EdgeIDType eid = eids_data[i];
+            PPR_list[src] = compute_s1_s2(src, dst, index0, eid, ts);
+            norm_list[src] = norm_list[src]*beta+beta;
+            if(src!=dst){
+                PPR_list[dst] = compute_s1_s2(dst, src, index0, eid, ts);
+                norm_list[dst] = norm_list[dst]*beta+beta;                
+            }
+        }
+    }
+    this->val_norm_list.assign(this->norm_list.begin(), this->norm_list.end());
+    this->val_PPR_list.assign(this->PPR_list.begin(), this->PPR_list.end());
+}
\ No newline at end of file

docs/source/guide/install_guide.rst

 Installation
 ============
 
-.. note::
-    安装方法
\ No newline at end of file
+This document provides step-by-step instructions for deploying and installing the StarryGL. Follow the steps below to ensure a successful installation.
+
+Prerequisites
+=============
+
+Before starting the installation, make sure you meet the following prerequisites:
+
+- Operating System:  Only Support for Linux.
+- Python Version: Requires Python 3.10.
+- CUDA Toolkit Version: Requires CUDA Toolkit 11.8.
+
+Installation Steps
+==================
+
+1. Clone the StarryGL Git Repository
+-------------------------------------
+
+First, clone the StarryGL Git repository, including the submodule dependencies. Run the following command to clone the repository:
+
+.. code-block:: bash
+
+    $ git clone --recursive http://115.236.33.124:7001/wjie98/starrygl.git
+
+If you have already cloned the repository without the submodules, you can run the following command to download the submodules:
+
+.. code-block:: bash
+
+    $ git submodule update --init --recursive
+
+This will ensure that all the necessary submodule dependencies are downloaded.
+
+2. Install Dependencies
+-----------------------
+
+Before installation, ensure that the necessary dependencies are installed on your system. Run the following command to install these dependencies:
+
+.. code-block:: bash
+
+    $ pip install -r requirements.txt
+
+3. Execute the Installation Script
+----------------------------------
+
+Locate the installation script file `install.sh` in the installation directory and execute the following command to install StarryGL:
+
+.. code-block:: bash
+
+    $ bash ./install.sh
+
+4. Install StarryGL
+------------------------
+
+The installation script will execute the `python setup.py install` command to install StarryGL. Make sure you are in the installation directory of StarryGL and execute the following command:
+
+.. code-block:: bash
+
+    $ python setup.py install
+
+5. Verify the Installation
+--------------------------
+
+After the installation, you can verify if StarryGL is successfully installed by executing the following command:
+
+.. code-block:: bash
+
+    $ python -c "import starrygl; print(starrygl.__version__)"
+
+If the StarryGL version is successfully displayed, it means the installation was successful.

docs/source/index.rst

@@ -2,6 +2,9 @@ StarryGL Documentation
 ======================
 
 .. toctree::
+   :maxdepth: 1
+   :caption: Contents
+   
    guide/index
    tutorial/index
    advanced/index

install.sh

@@ -11,3 +11,8 @@ cmake .. \
 && mkdir ../starrygl/lib \
 && cp lib*.so ../starrygl/lib/ \
 && patchelf --set-rpath '$ORIGIN:$ORIGIN/lib' --force-rpath ../starrygl/lib/*.so
+
+
+    # -DCMAKE_PREFIX_PATH="/home/zlj/.miniconda3/envs/dgnn/lib/python3.10/site-packages" \
+    # -DPython3_ROOT_DIR="/home/zlj/.miniconda3/envs/dgnn" \
+    # -DCUDA_TOOLKIT_ROOT_DIR="/home/zlj/local/cuda-12.2" \
\ No newline at end of file

starrygl/module/modules.py

 import torch
-import dgl
 from os.path import abspath, join, dirname
 import sys
 sys.path.insert(0, join(abspath(dirname(__file__))))

starrygl/sample/data_loader.py

@@ -155,10 +155,10 @@ class DistributedDataLoader:
         self.expected_idx = data_size // self.batch_size if self.drop_last is True else int(math.ceil(data_size/self.batch_size))
 
         if dist.get_world_size() > 1:
-            num_epochs = torch.tensor([self.expected_idx],dtype = torch.long,device=self.device) 
-            print(num_epochs)
-            dist.all_reduce(num_epochs, op=op)
-            self.expected_idx = int(num_epochs.item())
+            num_batchs = torch.tensor([self.expected_idx],dtype = torch.long,device=self.device) 
+            print("num_batchs:", num_batchs)
+            dist.all_reduce(num_batchs, op=op)
+            self.expected_idx = int(num_batchs.item())
 
     def _next_data(self):   
         if self.current_pos >= self.dataset.len:

starrygl/sample/sample_core/Utils copy.py

-import os.path as osp
-import torch
-class GraphData():
-    def __init__(self, path):
-        assert path is not None and osp.exists(path),'path 不存在'
-        id,edge_index,data,partptr =torch.load(path)
-        # 当前分区序号
-        self.partition_id = id
-        # 总分区数
-        self.partitions = partptr.numel() - 1
-        # 全图结构数据
-        self.num_nodes = partptr[self.partitions]
-        self.num_edges = edge_index[0].numel()
-        self.edge_index = edge_index
-        # 该分区下的数据（包含特征向量和子图结构）pyg Data数据结构
-        self.data = data
-        # 分区映射关系
-        self.partptr = partptr
-        self.eid = [i for i in range(self.num_edges)]
-
-    def __init__(self, id, edge_index, data, partptr, timestamp=None):
-        # 当前分区序号
-        self.partition_id = id
-        # 总分区数
-        self.partitions = partptr.numel() - 1
-        # 全图结构数据
-        self.num_nodes = partptr[self.partitions]
-        if edge_index is not None:
-            self.num_edges = edge_index[0].numel()
-        self.edge_index = edge_index
-        self.edge_ts = timestamp
-        # 该分区下的数据（包含特征向量和子图结构）pyg Data数据结构
-        self.data = data
-        # 分区映射关系
-        self.partptr = partptr
-        # edge id
-        self.eid = torch.tensor([i for i in range(0, self.num_edges)])
-
-    def select_attr(self,index):
-        return torch.index_select(self.data.x,0,index)
-
-    #返回全局的节点id 所对应的分区
-    def get_part_num(self):
-        return self.data.x.size()[0]
-
-    def select_attr(self,index):
-        return torch.index_select(self.data.x,0,index)
-    def select_y(self,index):
-        return torch.index_select(self.data.y,0,index)
-    #返回全局的节点id 所对应的分区
-    def get_localId_by_partitionId(self,id,index):
-        #print(index)
-        if(id == -1 or id == 0):
-            return index
-        else:
-            return torch.add(index,-self.partptr[id])
-    def get_globalId_by_partitionId(self,id,index):
-        if(id == -1 or id == 0):
-            return index
-        else:
-            return torch.add(index,self.partptr[id])
-
-    def get_node_num(self):    
-        return self.num_nodes
-    
-    
-    def localId_to_globalId(self,id,partitionId:int = -1):
-        '''
-        将分区partitionId内的点id映射为全局的id
-        '''
-        if partitionId == -1:
-            partitionId = self.partition_id
-        assert id >=self.partptr[self.partition_id] and id < self.partptr[self.partition_id+1]
-        ids_before = 0
-        if self.partition_id>0:
-            ids_before = self.partptr[self.partition_id-1]
-        return id+ids_before
-    
-    def get_partitionId_by_globalId(self,id):
-        '''
-        通过全局id得到对应的分区序号
-        '''
-        partitionId = -1
-        assert id>=0 and id<self.num_nodes,'id 超过范围'
-        for i in range(self.partitions):
-            if id>=self.partptr[i] and id<self.partptr[i+1]:
-                partitionId = i
-                break
-        assert partitionId>=0, 'id 不存在对应的分区'
-        return partitionId
-    
-    def get_nodes_by_partitionId(self,id):
-        '''
-        根据partitioId 返回该分区的节点数量
-        
-        '''
-        assert id>=0 and id<self.partitions,'partitionId 非法'
-        return (int)(self.partptr[id+1]-self.partptr[id])
-        
-    
-    def __repr__(self):
-        return (f'{self.__class__.__name__}(\n'
-                f'  partition_id={self.partition_id}\n'
-                f'  data={self.data},\n'
-                f'  global_info('
-                f'num_nodes={self.num_nodes},'
-                f' num_edges={self.num_edges},'
-                f' num_parts={self.partitions},'
-                f' edge_index=[2,{self.edge_index[0].numel()}])\n'
-                f')')

starrygl/sample/sample_core/neighbor_sampler.py

@@ -8,7 +8,7 @@ import torch.multiprocessing as mp
 from typing import Optional, Tuple
 
 from .base import BaseSampler, NegativeSampling, SampleOutput, SampleType
-# from sample_cores import ParallelSampler, get_neighbors, heads_unique
+from starrygl.lib.libstarrygl_sampler import ParallelSampler, get_neighbors, heads_unique
 
 from torch.distributed.rpc import rpc_async
 

starrygl/sample/sample_core/reddit_demo_before.py

@@ -62,7 +62,7 @@ def test():
     torch.backends.cudnn.benchmark = False
     torch.backends.cudnn.deterministic = True
 
-    g_data, df = load_gdelt_dataset()
+    g_data, df = load_reddit_dataset()
     print(g_data)
     # for worker in [1,2,3,4,5,6,7,8,9,10,20,30]:
     # import random

starrygl/sample/sample_core/reddit_demo_static.py

@@ -70,7 +70,7 @@ def test():
     torch.backends.cudnn.benchmark = False
     torch.backends.cudnn.deterministic = True
 
-    g_data = load_ogb_dataset()
+    g_data = load_reddit_dataset()
     print(g_data)
 
     from .neighbor_sampler import NeighborSampler, get_neighbors

train_tgnn.py

@@ -43,7 +43,7 @@ parser.add_argument('--rank', default=0, type=int, metavar='W',
 parser.add_argument('--patience', type=int, default=5, help='Patience for early stopping')
 parser.add_argument('--world_size', default=1, type=int, metavar='W',
                     help='number of negative samples')
-parser.add_argument('--dataname', default=1, type=str, metavar='W',
+parser.add_argument('--dataname', default="WIKI", type=str, metavar='W',
                     help='name of dataset')
 parser.add_argument('--model', default='TGN', type=str, metavar='W',
                     help='name of model')
@@ -52,14 +52,15 @@ from sklearn.metrics import average_precision_score, roc_auc_score
 import torch
 import time
 import random
-import dgl
 import numpy as np
 from sklearn.metrics import average_precision_score, roc_auc_score
 from torch.nn.parallel import DistributedDataParallel as DDP
-#os.environ['CUDA_VISIBLE_DEVICES'] = str(args.rank)
-#os.environ["RANK"] = str(args.rank)
-#os.environ["WORLD_SIZE"] = str(args.world_size)
-#os.environ["LOCAL_RANK"] = str(0)
+
+os.environ['CUDA_VISIBLE_DEVICES'] = str(args.rank)
+os.environ["RANK"] = str(args.rank)
+os.environ["WORLD_SIZE"] = str(args.world_size)
+os.environ["LOCAL_RANK"] = str(0)
+
 torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
 os.environ["MASTER_ADDR"] = '127.0.0.1'
 os.environ["MASTER_PORT"] = '9437'
@@ -152,7 +153,7 @@ def main():
     else:
         model = GeneralModel(gnn_dim_node, gnn_dim_edge, sample_param, memory_param, gnn_param, train_param)
         device = torch.device('cpu')
-    model = DDP(model,find_unused_parameters=True)
+    model = DDP(model,find_unused_parameters=False)
     train_stream = torch.cuda.Stream()
     send_stream = torch.cuda.Stream()
     scatter_stream = torch.cuda.Stream()
@@ -223,8 +224,11 @@ def main():
         auc_mrr = torch.empty([loader.expected_idx*world_size],dtype = torch.float,device = 'cuda')
         dist.all_gather_into_tensor(apc,torch.tensor(aps,device ='cuda',dtype=torch.float))
         dist.all_gather_into_tensor(auc_mrr,torch.tensor(aucs_mrrs,device ='cuda',dtype=torch.float))
-        ap = float(torch.tensor(apc).mean())
-        auc_mrr = float(torch.tensor(auc_mrr).mean())
+        # ap = float(torch.tensor(apc).mean())
+        # auc_mrr = float(torch.tensor(auc_mrr).mean())
+        ap = float(apc.clone().mean())
+        auc_mrr = float(auc_mrr.clone().mean())
+
         return ap, auc_mrr
     creterion = torch.nn.BCEWithLogitsLoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=train_param['lr'])