Skip to content

GitLab

Unverified Commit 80b99adb authored by Matthias Fey's avatar Matthias Fey Committed by GitHub
Browse files

Merge pull request #57 from rusty1s/wheel 

[WIP] Python wheels
parents 0194ebb6 bc476876
Hide whitespace changes
Inline Side-by-side
Showing with 422 additions and 315 deletions
.coveragerc
View file @ 80b99adb
......@@ -3,5 +3,7 @@ source=torch_cluster
[report]
exclude_lines =
pragma: no cover
cuda
torch.jit.script
raise
except
is_cuda
.travis.yml
View file @ 80b99adb
language: shell
os:
- linux
- osx
- windows
env:
global:
- CUDA_HOME=/usr/local/cuda
jobs:
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cpu
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu92
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu100
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu101
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cpu
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu92
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu100
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu101
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cpu
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu92
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu100
- TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu101
jobs:
include:
- os: linux
language: python
python: 3.7
addons:
apt:
sources:
- ubuntu-toolchain-r-test
packages:
- gcc-5
- g++-5
env:
- CC=gcc-5
- CXX=g++-5
- os: osx
language: sh
before_cache:
- brew cleanup
cache:
directories:
- $HOME/Library/Caches/Homebrew
- /usr/local/Homebrew
addons:
homebrew:
packages: python3
before_install:
- python3 -m pip install --upgrade virtualenv
- virtualenv -p python3 --system-site-packages "$HOME/venv"
- source "$HOME/venv/bin/activate"
env:
- CC=clang
- CXX=clang++
exclude: # Exclude *all* macOS CUDA jobs and Windows CUDA 9.2/10.0 jobs.
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu92
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu100
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu101
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu92
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu100
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu101
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu92
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu100
- os: osx
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu101
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu92
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu100
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu92
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.7 IDX=cu100
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu92
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.6 IDX=cu100
- os: windows
env: TORCH_VERSION=1.4.0 PYTHON_VERSION=3.8 IDX=cu101
install:
- pip install numpy
- pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
- pip install pycodestyle
- pip install flake8
- pip install codecov
- source script/cuda.sh
- source script/conda.sh
- conda create --yes -n test python="${PYTHON_VERSION}"
- source activate test
- conda install pytorch=${TORCH_VERSION} ${TOOLKIT} -c pytorch --yes
- source script/torch.sh
- pip install flake8 codecov
- python setup.py install
script:
- python -c "import torch; print(torch.__version__)"
- pycodestyle .
- flake8 .
- python setup.py install
- python setup.py test
after_success:
- python setup.py bdist_wheel --dist-dir=dist/torch-${TORCH_VERSION}
- python script/rename_wheel.py ${IDX}
- codecov
deploy:
provider: s3
region: eu-central-1
edge: true
access_key_id: ${S3_ACCESS_KEY}
secret_access_key: ${S3_SECRET_ACCESS_KEY}
bucket: pytorch-geometric.com
local_dir: dist/torch-${TORCH_VERSION}
upload_dir: whl/torch-${TORCH_VERSION}
acl: public_read
on:
repo: rusty1s/pytorch_cluster
tags: true
notifications:
email: false
LICENSE
View file @ 80b99adb
Copyright (c) 2019 Matthias Fey <matthias.fey@tu-dortmund.de>
Copyright (c) 2020 Matthias Fey <matthias.fey@tu-dortmund.de>
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
......
MANIFEST.in
View file @ 80b99adb
include README.md
include LICENSE
recursive-include cpu *
recursive-include cuda *
recursive-exclude test *
recursive-include csrc *
README.md
View file @ 80b99adb
......@@ -27,7 +27,30 @@ All included operations work on varying data types and are implemented both for
## Installation
Ensure that at least PyTorch 1.1.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*:
### Binaries
We provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://s3.eu-central-1.amazonaws.com/pytorch-geometric.com/whl/index.html).
To install from binaries, simply run
```
pip install torch-cluster==latest+${CUDA} -f https://s3.eu-central-1.amazonaws.com/pytorch-geometric.com/whl/torch-1.4.0.html
```
where `${CUDA}` should be replaced by either `cpu`, `cu92`, `cu100` or `cu101` depending on your PyTorch installation.
| | `cpu` | `cu92` | `cu100` | `cu101` |
|-------------|-------|--------|---------|---------|
| **Linux** | ✅ | ✅ | ✅ | ✅ |
| **Windows** | ✅ | ❌ | ❌ | ✅ |
| **macOS** | ✅ | | | |
### From source
Ensure that at least PyTorch 1.4.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*:
```
$ python -c "import torch; print(torch.__version__)"
>>> 1.4.0
```
$ python -c "import torch; print(torch.__version__)"
......@@ -46,10 +69,16 @@ Then run:
pip install torch-cluster
```
If you are running into any installation problems, please create an [issue](https://github.com/rusty1s/pytorch_cluster/issues).
Be sure to import `torch` first before using this package to resolve symbols the dynamic linker must see.
When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:
```
export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
```
## Functions
## Graclus
### Graclus
A greedy clustering algorithm of picking an unmarked vertex and matching it with one its unmarked neighbors (that maximizes its edge weight).
The GPU algorithm is adapted from Fagginger Auer and Bisseling: [A GPU Algorithm for Greedy Graph Matching](http://www.staff.science.uu.nl/~bisse101/Articles/match12.pdf) (LNCS 2012)
......@@ -70,7 +99,7 @@ print(cluster)
tensor([0, 0, 1])
```
## VoxelGrid
### VoxelGrid
A clustering algorithm, which overlays a regular grid of user-defined size over a point cloud and clusters all points within a voxel.
......@@ -89,7 +118,7 @@ print(cluster)
tensor([0, 5, 3, 0, 1])
```
## FarthestPointSampling
### FarthestPointSampling
A sampling algorithm, which iteratively samples the most distant point with regard to the rest points.
......@@ -107,7 +136,7 @@ print(sample)
tensor([0, 3])
```
## kNN-Graph
### kNN-Graph
Computes graph edges to the nearest *k* points.
......@@ -126,7 +155,7 @@ tensor([[1, 2, 0, 3, 0, 3, 1, 2],
[0, 0, 1, 1, 2, 2, 3, 3]])
```
## Radius-Graph
### Radius-Graph
Computes graph edges to all points within a given distance.
......@@ -145,7 +174,7 @@ tensor([[1, 2, 0, 3, 0, 3, 1, 2],
[0, 0, 1, 1, 2, 2, 3, 3]])
```
## Nearest
### Nearest
Clusters points in *x* together which are nearest to a given query point in *y*.
......@@ -165,7 +194,7 @@ print(cluster)
tensor([0, 0, 1, 1])
```
## RandomWalk-Sampling
### RandomWalk-Sampling
Samples random walks of length `walk_length` from all node indices in `start` in the graph given by `(row, col)`.
......
cpu/compat.h deleted 100644 → 0
View file @ 0194ebb6
#ifdef VERSION_GE_1_3
#define DATA_PTR data_ptr
#else
#define DATA_PTR data
#endif
cpu/fps.cpp deleted 100644 → 0
View file @ 0194ebb6
#include <torch/extension.h>
#include "compat.h"
#include "utils.h"
at::Tensor get_dist(at::Tensor x, ptrdiff_t index) {
return (x - x[index]).norm(2, 1);
}
at::Tensor fps(at::Tensor x, at::Tensor batch, float ratio, bool random) {
auto batch_size = batch[-1].DATA_PTR<int64_t>()[0] + 1;
auto deg = degree(batch, batch_size);
auto cum_deg = at::cat({at::zeros(1, deg.options()), deg.cumsum(0)}, 0);
auto k = (deg.toType(at::kFloat) * ratio).ceil().toType(at::kLong);
auto cum_k = at::cat({at::zeros(1, k.options()), k.cumsum(0)}, 0);
auto out = at::empty(cum_k[-1].DATA_PTR<int64_t>()[0], batch.options());
auto cum_deg_d = cum_deg.DATA_PTR<int64_t>();
auto k_d = k.DATA_PTR<int64_t>();
auto cum_k_d = cum_k.DATA_PTR<int64_t>();
auto out_d = out.DATA_PTR<int64_t>();
for (ptrdiff_t b = 0; b < batch_size; b++) {
auto index = at::range(cum_deg_d[b], cum_deg_d[b + 1] - 1, out.options());
auto y = x.index_select(0, index);
ptrdiff_t start = 0;
if (random) {
start = at::randperm(y.size(0), batch.options()).DATA_PTR<int64_t>()[0];
}
out_d[cum_k_d[b]] = cum_deg_d[b] + start;
auto dist = get_dist(y, start);
for (ptrdiff_t i = 1; i < k_d[b]; i++) {
ptrdiff_t argmax = dist.argmax().DATA_PTR<int64_t>()[0];
out_d[cum_k_d[b] + i] = cum_deg_d[b] + argmax;
dist = at::min(dist, get_dist(y, argmax));
}
}
return out;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fps", &fps, "Farthest Point Sampling (CPU)");
}
cpu/graclus.cpp deleted 100644 → 0
View file @ 0194ebb6
#include <torch/extension.h>
#include "compat.h"
#include "utils.h"
at::Tensor graclus(at::Tensor row, at::Tensor col, int64_t num_nodes) {
std::tie(row, col) = remove_self_loops(row, col);
std::tie(row, col) = rand(row, col);
std::tie(row, col) = to_csr(row, col, num_nodes);
auto row_data = row.DATA_PTR<int64_t>(), col_data = col.DATA_PTR<int64_t>();
auto perm = at::randperm(num_nodes, row.options());
auto perm_data = perm.DATA_PTR<int64_t>();
auto cluster = at::full(num_nodes, -1, row.options());
auto cluster_data = cluster.DATA_PTR<int64_t>();
for (int64_t i = 0; i < num_nodes; i++) {
auto u = perm_data[i];
if (cluster_data[u] >= 0)
continue;
cluster_data[u] = u;
for (int64_t j = row_data[u]; j < row_data[u + 1]; j++) {
auto v = col_data[j];
if (cluster_data[v] >= 0)
continue;
cluster_data[u] = std::min(u, v);
cluster_data[v] = std::min(u, v);
break;
}
}
return cluster;
}
at::Tensor weighted_graclus(at::Tensor row, at::Tensor col, at::Tensor weight,
int64_t num_nodes) {
std::tie(row, col, weight) = remove_self_loops(row, col, weight);
std::tie(row, col, weight) = to_csr(row, col, weight, num_nodes);
auto row_data = row.DATA_PTR<int64_t>(), col_data = col.DATA_PTR<int64_t>();
auto perm = at::randperm(num_nodes, row.options());
auto perm_data = perm.DATA_PTR<int64_t>();
auto cluster = at::full(num_nodes, -1, row.options());
auto cluster_data = cluster.DATA_PTR<int64_t>();
AT_DISPATCH_ALL_TYPES(weight.scalar_type(), "weighted_graclus", [&] {
auto weight_data = weight.DATA_PTR<scalar_t>();
for (int64_t i = 0; i < num_nodes; i++) {
auto u = perm_data[i];
if (cluster_data[u] >= 0)
continue;
int64_t v_max = u;
scalar_t w_max = 0;
for (int64_t j = row_data[u]; j < row_data[u + 1]; j++) {
auto v = col_data[j];
if (cluster_data[v] >= 0)
continue;
if (weight_data[j] >= w_max) {
v_max = v;
w_max = weight_data[j];
}
}
cluster_data[u] = std::min(u, v_max);
cluster_data[v_max] = std::min(u, v_max);
}
});
return cluster;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("graclus", &graclus, "Graclus (CPU)");
m.def("weighted_graclus", &weighted_graclus, "Weighted Graclus (CPU)");
}
cpu/grid.cpp deleted 100644 → 0
View file @ 0194ebb6
#include <torch/extension.h>
at::Tensor grid(at::Tensor pos, at::Tensor size, at::Tensor start,
at::Tensor end) {
pos = pos - start.view({1, -1});
auto num_voxels = ((end - start) / size).toType(at::kLong) + 1;
num_voxels = num_voxels.cumprod(0);
num_voxels = at::cat({at::ones(1, num_voxels.options()), num_voxels}, 0);
auto index = at::empty(size.size(0), num_voxels.options());
at::arange_out(index, size.size(0));
num_voxels = num_voxels.index_select(0, index);
auto cluster = (pos / size.view({1, -1})).toType(at::kLong);
cluster *= num_voxels.view({1, -1});
cluster = cluster.sum(1);
return cluster;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grid", &grid, "Grid (CPU)"); }
cpu/rw.cpp deleted 100644 → 0
View file @ 0194ebb6
#include <torch/extension.h>
#include "compat.h"
#include "utils.h"
at::Tensor rw(at::Tensor row, at::Tensor col, at::Tensor start,
size_t walk_length, float p, float q, size_t num_nodes) {
auto deg = degree(row, num_nodes);
auto cum_deg = at::cat({at::zeros(1, deg.options()), deg.cumsum(0)}, 0);
auto rand = at::rand({start.size(0), (int64_t)walk_length},
start.options().dtype(at::kFloat));
auto out =
at::full({start.size(0), (int64_t)walk_length + 1}, -1, start.options());
auto deg_d = deg.DATA_PTR<int64_t>();
auto cum_deg_d = cum_deg.DATA_PTR<int64_t>();
auto col_d = col.DATA_PTR<int64_t>();
auto start_d = start.DATA_PTR<int64_t>();
auto rand_d = rand.DATA_PTR<float>();
auto out_d = out.DATA_PTR<int64_t>();
for (ptrdiff_t n = 0; n < start.size(0); n++) {
int64_t cur = start_d[n];
auto i = n * (walk_length + 1);
out_d[i] = cur;
for (ptrdiff_t l = 1; l <= (int64_t)walk_length; l++) {
cur = col_d[cum_deg_d[cur] +
int64_t(rand_d[n * walk_length + (l - 1)] * deg_d[cur])];
out_d[i + l] = cur;
}
}
return out;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rw", &rw, "Random Walk Sampling (CPU)");
}
cpu/utils.h deleted 100644 → 0
View file @ 0194ebb6
#pragma once
#include <torch/extension.h>
std::tuple<at::Tensor, at::Tensor> remove_self_loops(at::Tensor row,
at::Tensor col) {
auto mask = row != col;
return std::make_tuple(row.masked_select(mask), col.masked_select(mask));
}
std::tuple<at::Tensor, at::Tensor, at::Tensor>
remove_self_loops(at::Tensor row, at::Tensor col, at::Tensor weight) {
auto mask = row != col;
return std::make_tuple(row.masked_select(mask), col.masked_select(mask),
weight.masked_select(mask));
}
std::tuple<at::Tensor, at::Tensor> rand(at::Tensor row, at::Tensor col) {
auto perm = at::randperm(row.size(0), row.options());
return std::make_tuple(row.index_select(0, perm), col.index_select(0, perm));
}
std::tuple<at::Tensor, at::Tensor> sort_by_row(at::Tensor row, at::Tensor col) {
at::Tensor perm;
std::tie(row, perm) = row.sort();
return std::make_tuple(row, col.index_select(0, perm));
}
std::tuple<at::Tensor, at::Tensor, at::Tensor>
sort_by_row(at::Tensor row, at::Tensor col, at::Tensor weight) {
at::Tensor perm;
std::tie(row, perm) = row.sort();
return std::make_tuple(row, col.index_select(0, perm),
weight.index_select(0, perm));
}
at::Tensor degree(at::Tensor row, int64_t num_nodes) {
auto zero = at::zeros(num_nodes, row.options());
auto one = at::ones(row.size(0), row.options());
return zero.scatter_add_(0, row, one);
}
std::tuple<at::Tensor, at::Tensor> to_csr(at::Tensor row, at::Tensor col,
int64_t num_nodes) {
std::tie(row, col) = sort_by_row(row, col);
row = degree(row, num_nodes).cumsum(0);
row = at::cat({at::zeros(1, row.options()), row}, 0); // Prepend zero.
return std::make_tuple(row, col);
}
std::tuple<at::Tensor, at::Tensor, at::Tensor>
to_csr(at::Tensor row, at::Tensor col, at::Tensor weight, int64_t num_nodes) {
std::tie(row, col, weight) = sort_by_row(row, col, weight);
row = degree(row, num_nodes).cumsum(0);
row = at::cat({at::zeros(1, row.options()), row}, 0); // Prepend zero.
return std::make_tuple(row, col, weight);
}
csrc/cpu/fps_cpu.cpp 0 → 100644
View file @ 80b99adb
#include "fps_cpu.h"
#include "utils.h"
inline torch::Tensor get_dist(torch::Tensor x, int64_t idx) {
return (x - x[idx]).norm(2, 1);
}
torch::Tensor fps_cpu(torch::Tensor src, torch::Tensor ptr, double ratio,
bool random_start) {
CHECK_CPU(src);
CHECK_CPU(ptr);
CHECK_INPUT(ptr.dim() == 1);
AT_ASSERTM(ratio > 0 && ratio < 1, "Invalid input");
src = src.view({src.size(0), -1}).contiguous();
ptr = ptr.contiguous();
auto batch_size = ptr.size(0) - 1;
auto deg = ptr.narrow(0, 1, batch_size) - ptr.narrow(0, 0, batch_size);
auto out_ptr = deg.toType(torch::kFloat) * (float)ratio;
out_ptr = out_ptr.ceil().toType(torch::kLong).cumsum(0);
auto out = torch::empty(out_ptr[-1].data_ptr<int64_t>()[0], ptr.options());
auto ptr_data = ptr.data_ptr<int64_t>();
auto out_ptr_data = out_ptr.data_ptr<int64_t>();
auto out_data = out.data_ptr<int64_t>();
int64_t src_start = 0, out_start = 0, src_end, out_end;
for (auto b = 0; b < batch_size; b++) {
src_end = ptr_data[b + 1], out_end = out_ptr_data[b];
auto y = src.narrow(0, src_start, src_end - src_start);
int64_t start_idx = 0;
if (random_start) {
start_idx = rand() % src.size(0);
}
out_data[out_start] = src_start + start_idx;
auto dist = get_dist(y, start_idx);
for (auto i = 1; i < out_end - out_start; i++) {
int64_t argmax = dist.argmax().data_ptr<int64_t>()[0];
out_data[out_start + i] = src_start + argmax;
dist = torch::min(dist, get_dist(y, argmax));
}
src_start = src_end, out_start = out_end;
}
return out;
}
csrc/cpu/fps_cpu.h 0 → 100644
View file @ 80b99adb
#pragma once
#include <torch/extension.h>
torch::Tensor fps_cpu(torch::Tensor src, torch::Tensor ptr, double ratio,
bool random_start);
csrc/cpu/graclus_cpu.cpp 0 → 100644
View file @ 80b99adb
#include "graclus_cpu.h"
#include "utils.h"
torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
torch::optional<torch::Tensor> optional_weight) {
CHECK_CPU(rowptr);
CHECK_CPU(col);
CHECK_INPUT(rowptr.dim() == 1 && col.dim() == 1);
if (optional_weight.has_value()) {
CHECK_CPU(optional_weight.value());
CHECK_INPUT(optional_weight.value().dim() == 1);
CHECK_INPUT(optional_weight.value().numel() == col.numel());
}
int64_t num_nodes = rowptr.numel() - 1;
auto out = torch::full(num_nodes, -1, rowptr.options());
auto node_perm = torch::randperm(num_nodes, rowptr.options());
auto rowptr_data = rowptr.data_ptr<int64_t>();
auto col_data = col.data_ptr<int64_t>();
auto node_perm_data = node_perm.data_ptr<int64_t>();
auto out_data = out.data_ptr<int64_t>();
if (!optional_weight.has_value()) {
for (int64_t n = 0; n < num_nodes; n++) {
auto u = node_perm_data[n];
if (out_data[u] >= 0)
continue;
out_data[u] = u;
int64_t row_start = rowptr_data[u], row_end = rowptr_data[u + 1];
for (auto e = 0; e < row_end - row_start; e++) {
auto v = col_data[row_start + e];
if (out_data[v] >= 0)
continue;
out_data[u] = std::min(u, v);
out_data[v] = std::min(u, v);
break;
}
}
} else {
auto weight = optional_weight.value();
AT_DISPATCH_ALL_TYPES(weight.scalar_type(), "weighted_graclus", [&] {
auto weight_data = weight.data_ptr<scalar_t>();
for (auto n = 0; n < num_nodes; n++) {
auto u = node_perm_data[n];
if (out_data[u] >= 0)
continue;
auto v_max = u;
scalar_t w_max = (scalar_t)0.;
for (auto e = rowptr_data[u]; e < rowptr_data[u + 1]; e++) {
auto v = col_data[e];
if (out_data[v] >= 0)
continue;
if (weight_data[e] >= w_max) {
v_max = v;
w_max = weight_data[e];
}
}
out_data[u] = std::min(u, v_max);
out_data[v_max] = std::min(u, v_max);
}
});
}
return out;
}
csrc/cpu/graclus_cpu.h 0 → 100644
View file @ 80b99adb
#pragma once
#include <torch/extension.h>
torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
torch::optional<torch::Tensor> optional_weight);
csrc/cpu/grid_cpu.cpp 0 → 100644
View file @ 80b99adb
#include "grid_cpu.h"
#include "utils.h"
torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
torch::optional<torch::Tensor> optional_start,
torch::optional<torch::Tensor> optional_end) {
CHECK_CPU(pos);
CHECK_CPU(size);
if (optional_start.has_value())
CHECK_CPU(optional_start.value());
if (optional_start.has_value())
CHECK_CPU(optional_start.value());
pos = pos.view({pos.size(0), -1});
CHECK_INPUT(size.numel() == pos.size(1));
if (!optional_start.has_value())
optional_start = std::get<0>(pos.min(0));
else
CHECK_INPUT(optional_start.value().numel() == pos.size(1));
if (!optional_end.has_value())
optional_end = std::get<0>(pos.max(0));
else
CHECK_INPUT(optional_start.value().numel() == pos.size(1));
auto start = optional_start.value();
auto end = optional_end.value();
pos = pos - start.unsqueeze(0);
auto num_voxels = ((end - start) / size).toType(torch::kLong) + 1;
num_voxels = num_voxels.cumprod(0);
num_voxels =
torch::cat({torch::ones(1, num_voxels.options()), num_voxels}, 0);
num_voxels = num_voxels.narrow(0, 0, size.size(0));
auto out = (pos / size.view({1, -1})).toType(torch::kLong);
out *= num_voxels.view({1, -1});
out = out.sum(1);
return out;
}
csrc/cpu/grid_cpu.h 0 → 100644
View file @ 80b99adb
#pragma once
#include <torch/extension.h>
torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
torch::optional<torch::Tensor> optional_start,
torch::optional<torch::Tensor> optional_end);
csrc/cpu/rw_cpu.cpp 0 → 100644
View file @ 80b99adb
#include "rw_cpu.h"
#include "utils.h"
torch::Tensor random_walk_cpu(torch::Tensor rowptr, torch::Tensor col,
torch::Tensor start, int64_t walk_length,
double p, double q) {
CHECK_CPU(rowptr);
CHECK_CPU(col);
CHECK_CPU(start);
CHECK_INPUT(rowptr.dim() == 1);
CHECK_INPUT(col.dim() == 1);
CHECK_INPUT(start.dim() == 1);
auto rand = torch::rand({start.size(0), walk_length},
start.options().dtype(torch::kFloat));
auto out = torch::full({start.size(0), walk_length + 1}, -1, start.options());
auto rowptr_data = rowptr.data_ptr<int64_t>();
auto col_data = col.data_ptr<int64_t>();
auto start_data = start.data_ptr<int64_t>();
auto rand_data = rand.data_ptr<float>();
auto out_data = out.data_ptr<int64_t>();
for (auto n = 0; n < start.size(0); n++) {
auto cur = start_data[n];
auto offset = n * (walk_length + 1);
out_data[offset] = cur;
int64_t row_start, row_end;
for (auto l = 1; l <= walk_length; l++) {
row_start = rowptr_data[cur], row_end = rowptr_data[cur + 1];
cur = col_data[row_start + int64_t(rand_data[n * walk_length + (l - 1)] *
(row_end - row_start))];
out_data[offset + l] = cur;
}
}
return out;
}
csrc/cpu/rw_cpu.h 0 → 100644
View file @ 80b99adb
#pragma once
#include <torch/extension.h>
torch::Tensor random_walk_cpu(torch::Tensor rowptr, torch::Tensor col,
torch::Tensor start, int64_t walk_length,
double p, double q);
cpu/sampler.cpp → csrc/cpu/sampler_cpu.cpp
View file @ 80b99adb
#include <torch/extension.h>
#include "sampler_cpu.h"
#include "compat.h"
#include "utils.h"
at::Tensor neighbor_sampler(at::Tensor start, at::Tensor cumdeg, size_t size,
float factor) {
torch::Tensor neighbor_sampler_cpu(torch::Tensor start, torch::Tensor rowptr,
int64_t count, double factor) {
auto start_ptr = start.DATA_PTR<int64_t>();
auto cumdeg_ptr = cumdeg.DATA_PTR<int64_t>();
auto start_data = start.data_ptr<int64_t>();
auto rowptr_data = rowptr.data_ptr<int64_t>();
std::vector<int64_t> e_ids;
for (ptrdiff_t i = 0; i < start.size(0); i++) {
int64_t low = cumdeg_ptr[start_ptr[i]];
int64_t high = cumdeg_ptr[start_ptr[i] + 1];
size_t num_neighbors = high - low;
size_t size_i = size_t(ceil(factor * float(num_neighbors)));
size_i = (size_i < size) ? size_i : size;
for (auto i = 0; i < start.size(0); i++) {
auto row_start = rowptr_data[start_data[i]];
auto row_end = rowptr_data[start_data[i] + 1];
auto num_neighbors = row_end - row_start;
int64_t size = count;
if (count < 1) {
size = int64_t(ceil(factor * float(num_neighbors)));
}
// If the number of neighbors is approximately equal to the number of
// neighbors which are requested, we use `randperm` to sample without
// replacement, otherwise we sample random numbers into a set as long as
// necessary.
// replacement, otherwise we sample random numbers into a set as long
// as necessary.
std::unordered_set<int64_t> set;
if (size_i < 0.7 * float(num_neighbors)) {
while (set.size() < size_i) {
int64_t z = rand() % num_neighbors;
set.insert(z + low);
if (size < 0.7 * float(num_neighbors)) {
while (int64_t(set.size()) < size) {
int64_t sample = (rand() % num_neighbors) + row_start;
set.insert(sample);
}
std::vector<int64_t> v(set.begin(), set.end());
e_ids.insert(e_ids.end(), v.begin(), v.end());
} else {
auto sample = at::randperm(num_neighbors, start.options());
auto sample_ptr = sample.DATA_PTR<int64_t>();
for (size_t j = 0; j < size_i; j++) {
e_ids.push_back(sample_ptr[j] + low);
auto sample = at::randperm(num_neighbors, start.options()) + row_start;
auto sample_data = sample.data_ptr<int64_t>();
for (auto j = 0; j < size; j++) {
e_ids.push_back(sample_data[j]);
}
}
}
int64_t len = e_ids.size();
auto e_id = torch::from_blob(e_ids.data(), {len}, start.options()).clone();
return e_id;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("neighbor_sampler", &neighbor_sampler, "Neighbor Sampler (CPU)");
int64_t length = e_ids.size();
return torch::from_blob(e_ids.data(), {length}, start.options()).clone();
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment