p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!5624 Optimize the performance of GraphData.get_neg_sampled_neighbors 

Merge pull request !5624 from heleiwang/gnn_perf
This commit is contained in:
mindspore-ci-bot 2020-09-01 19:22:53 +08:00 committed by Gitee
parent 0a203c68bd 857cf2f77f
commit 8f3ebfd469
3 changed files with 30 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
mindspore/ccsrc/minddata/dataset/engine/gnn/graph_data_impl.cc
View File

@ -211,22 +211,22 @@ Status GraphDataImpl::GetSampledNeighbors(const std::vector<NodeIdType> &node_li
return Status::OK();
}
Status GraphDataImpl::NegativeSample(const std::vector<NodeIdType> &data,
const std::unordered_set<NodeIdType> &exclude_data, int32_t samples_num,
std::vector<NodeIdType> *out_samples) {
Status GraphDataImpl::NegativeSample(const std::vector<NodeIdType> &data, const std::vector<NodeIdType> shuffled_ids,
size_t *start_index, const std::unordered_set<NodeIdType> &exclude_data,
int32_t samples_num, std::vector<NodeIdType> *out_samples) {
CHECK_FAIL_RETURN_UNEXPECTED(!data.empty(), "Input data is empty.");
std::vector<NodeIdType> shuffled_id(data.size());
std::iota(shuffled_id.begin(), shuffled_id.end(), 0);
std::shuffle(shuffled_id.begin(), shuffled_id.end(), rnd_);
for (const auto &index : shuffled_id) {
if (exclude_data.find(data[index]) != exclude_data.end()) {
size_t index = *start_index;
for (size_t i = index; i < shuffled_ids.size(); ++i) {
++index;
if (exclude_data.find(data[shuffled_ids[i]]) != exclude_data.end()) {
continue;
}
out_samples->emplace_back(data[index]);
out_samples->emplace_back(data[shuffled_ids[i]]);
if (out_samples->size() >= samples_num) {
break;
}
}
*start_index = index;
return Status::OK();
}
@ -236,6 +236,13 @@ Status GraphDataImpl::GetNegSampledNeighbors(const std::vector<NodeIdType> &node
RETURN_IF_NOT_OK(CheckSamplesNum(samples_num));
RETURN_IF_NOT_OK(CheckNeighborType(neg_neighbor_type));
const std::vector<NodeIdType> &all_nodes = node_type_map_[neg_neighbor_type];
std::vector<NodeIdType> shuffled_id(all_nodes.size());
std::iota(shuffled_id.begin(), shuffled_id.end(), 0);
std::shuffle(shuffled_id.begin(), shuffled_id.end(), rnd_);
size_t start_index = 0;
bool need_shuffle = false;
std::vector<std::vector<NodeIdType>> neg_neighbors_vec;
neg_neighbors_vec.resize(node_list.size());
for (size_t node_idx = 0; node_idx < node_list.size(); ++node_idx) {
@ -247,12 +254,15 @@ Status GraphDataImpl::GetNegSampledNeighbors(const std::vector<NodeIdType> &node
std::transform(neighbors.begin(), neighbors.end(),
std::insert_iterator<std::unordered_set<NodeIdType>>(exclude_nodes, exclude_nodes.begin()),
[](const NodeIdType node) { return node; });
const std::vector<NodeIdType> &all_nodes = node_type_map_[neg_neighbor_type];
neg_neighbors_vec[node_idx].emplace_back(node->id());
if (all_nodes.size() > exclude_nodes.size()) {
while (neg_neighbors_vec[node_idx].size() < samples_num + 1) {
RETURN_IF_NOT_OK(NegativeSample(all_nodes, exclude_nodes, samples_num - neg_neighbors_vec[node_idx].size(),
RETURN_IF_NOT_OK(NegativeSample(all_nodes, shuffled_id, &start_index, exclude_nodes, samples_num + 1,
&neg_neighbors_vec[node_idx]));
if (start_index >= shuffled_id.size()) {
start_index = start_index % shuffled_id.size();
need_shuffle = true;
}
}
} else {
MS_LOG(DEBUG) << "There are no negative neighbors. node_id:" << node->id()
@ -262,6 +272,11 @@ Status GraphDataImpl::GetNegSampledNeighbors(const std::vector<NodeIdType> &node
neg_neighbors_vec[node_idx].emplace_back(kDefaultNodeId);
}
}
if (need_shuffle) {
std::shuffle(shuffled_id.begin(), shuffled_id.end(), rnd_);
start_index = 0;
need_shuffle = false;
}
}
RETURN_IF_NOT_OK(CreateTensorByVector<NodeIdType>(neg_neighbors_vec, DataType(DataType::DE_INT32), out));
return Status::OK();

5
mindspore/ccsrc/minddata/dataset/engine/gnn/graph_data_impl.h
View File

@ -251,8 +251,9 @@ class GraphDataImpl : public GraphData {
// @param int32_t samples_num -
// @param std::vector<NodeIdType> *out_samples - Sampling results returned
// @return Status - The error code return
Status NegativeSample(const std::vector<NodeIdType> &input_data, const std::unordered_set<NodeIdType> &exclude_data,
int32_t samples_num, std::vector<NodeIdType> *out_samples);
Status NegativeSample(const std::vector<NodeIdType> &data, const std::vector<NodeIdType> shuffled_ids,
size_t *start_index, const std::unordered_set<NodeIdType> &exclude_data, int32_t samples_num,
std::vector<NodeIdType> *out_samples);
Status CheckSamplesNum(NodeIdType samples_num);

2
tests/ut/python/dataset/test_graphdata_distributed.py
View File

@ -87,7 +87,7 @@ def test_graphdata_distributed():
p1 = Process(target=graphdata_startserver, args=(server_port,))
p1.start()
time.sleep(2)
time.sleep(5)
g = ds.GraphData(DATASET_FILE, 1, 'client', port=server_port)
nodes = g.get_all_nodes(1)