Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimization of biased sampling #270

Merged
merged 7 commits into from
Oct 31, 2023
Merged

Optimization of biased sampling #270

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
225785d
Added optimization of biased sampling
OlhaBabicheva Oct 31, 2023
80d5834
Small changes
OlhaBabicheva Oct 31, 2023
2722c17
Updated CHANGELOG.md
OlhaBabicheva Oct 31, 2023
2b7c2fb
Updated neighbor_kernel.cpp
OlhaBabicheva Oct 31, 2023
cf440e8
Added comments
OlhaBabicheva Oct 31, 2023
3fb34e9
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 31, 2023
633db5b
update
rusty1s Oct 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
### Added
### Changed
- Added `--biased` parameter to run benchmarks for biased sampling ([#267](https://github.com/pyg-team/pyg-lib/pull/267))
- Improved speed of biased sampling ([#270](https://github.com/pyg-team/pyg-lib/pull/270))
### Removed

## [0.3.0] - 2023-10-11
Expand Down
14 changes: 13 additions & 1 deletion pyg_lib/csrc/sampler/cpu/neighbor_kernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,19 @@ class NeighborSampler {

// Case 2: Multinomial sampling:
else {
const auto index = at::multinomial(weight, count, replace);
at::Tensor index;
if (replace) {
// at::multinomial only has good perfomance for `replace=true`, see:
// https://github.com/pytorch/pytorch/issues/11931
index = at::multinomial(weight, count, replace);
} else {
// For `replace=false`, we make use of the implementation of the
// "Weighted Random Sampling" paper:
// https://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf
const auto rand = at::empty_like(weight).uniform_();
const auto key = (rand.log() / weight);
rusty1s marked this conversation as resolved.
Show resolved Hide resolved
index = std::get<1>(key.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,
Expand Down