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 586 additions and 606 deletions
cuda/radius_kernel.cu deleted 100644 → 0
View file @ 0194ebb6
#include <ATen/ATen.h>
#include "compat.cuh"
#include "utils.cuh"
#define THREADS 1024
template <typename scalar_t>
__global__ void
radius_kernel(const scalar_t *__restrict__ x, const scalar_t *__restrict__ y,
const int64_t *__restrict__ batch_x,
const int64_t *__restrict__ batch_y, int64_t *__restrict__ row,
int64_t *__restrict__ col, scalar_t radius,
size_t max_num_neighbors, size_t dim) {
const ptrdiff_t batch_idx = blockIdx.x;
const ptrdiff_t idx = threadIdx.x;
const ptrdiff_t start_idx_x = batch_x[batch_idx];
const ptrdiff_t end_idx_x = batch_x[batch_idx + 1];
const ptrdiff_t start_idx_y = batch_y[batch_idx];
const ptrdiff_t end_idx_y = batch_y[batch_idx + 1];
for (ptrdiff_t n_y = start_idx_y + idx; n_y < end_idx_y; n_y += THREADS) {
size_t count = 0;
for (ptrdiff_t n_x = start_idx_x; n_x < end_idx_x; n_x++) {
scalar_t dist = 0;
for (ptrdiff_t d = 0; d < dim; d++) {
dist += (x[n_x * dim + d] - y[n_y * dim + d]) *
(x[n_x * dim + d] - y[n_y * dim + d]);
}
dist = sqrt(dist);
if (dist <= radius) {
row[n_y * max_num_neighbors + count] = n_y;
col[n_y * max_num_neighbors + count] = n_x;
count++;
}
if (count >= max_num_neighbors) {
break;
}
}
}
}
at::Tensor radius_cuda(at::Tensor x, at::Tensor y, float radius,
at::Tensor batch_x, at::Tensor batch_y,
size_t max_num_neighbors) {
cudaSetDevice(x.get_device());
auto batch_sizes = (int64_t *)malloc(sizeof(int64_t));
cudaMemcpy(batch_sizes, batch_x[-1].DATA_PTR<int64_t>(), sizeof(int64_t),
cudaMemcpyDeviceToHost);
auto batch_size = batch_sizes[0] + 1;
batch_x = degree(batch_x, batch_size);
batch_x = at::cat({at::zeros(1, batch_x.options()), batch_x.cumsum(0)}, 0);
batch_y = degree(batch_y, batch_size);
batch_y = at::cat({at::zeros(1, batch_y.options()), batch_y.cumsum(0)}, 0);
auto row = at::full(y.size(0) * max_num_neighbors, -1, batch_y.options());
auto col = at::full(y.size(0) * max_num_neighbors, -1, batch_y.options());
AT_DISPATCH_FLOATING_TYPES(x.scalar_type(), "radius_kernel", [&] {
radius_kernel<scalar_t><<<batch_size, THREADS>>>(
x.DATA_PTR<scalar_t>(), y.DATA_PTR<scalar_t>(),
batch_x.DATA_PTR<int64_t>(), batch_y.DATA_PTR<int64_t>(),
row.DATA_PTR<int64_t>(), col.DATA_PTR<int64_t>(), radius,
max_num_neighbors, x.size(1));
});
auto mask = row != -1;
return at::stack({row.masked_select(mask), col.masked_select(mask)}, 0);
}
cuda/response.cuh deleted 100644 → 0
View file @ 0194ebb6
#pragma once
#include <ATen/ATen.h>
#include "compat.cuh"
#define THREADS 1024
#define BLOCKS(N) (N + THREADS - 1) / THREADS
__global__ void respond_kernel(int64_t *__restrict__ cluster, int64_t *proposal,
int64_t *__restrict row,
int64_t *__restrict__ col, size_t numel) {
const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
const size_t stride = blockDim.x * gridDim.x;
for (int64_t u = index; u < numel; u += stride) {
if (cluster[u] != -2)
continue; // Only vist red nodes.
bool has_unmatched_neighbor = false;
for (int64_t i = row[u]; i < row[u + 1]; i++) {
auto v = col[i];
if (cluster[v] < 0)
has_unmatched_neighbor = true; // Unmatched neighbor found.
if (cluster[v] == -1 && proposal[v] == u) {
// Match first blue neighbhor v which proposed to u.
cluster[u] = min(u, v);
cluster[v] = min(u, v);
break;
}
}
if (!has_unmatched_neighbor)
cluster[u] = u;
}
}
void respond(at::Tensor cluster, at::Tensor proposal, at::Tensor row,
at::Tensor col) {
respond_kernel<<<BLOCKS(cluster.numel()), THREADS>>>(
cluster.DATA_PTR<int64_t>(), proposal.DATA_PTR<int64_t>(),
row.DATA_PTR<int64_t>(), col.DATA_PTR<int64_t>(), cluster.numel());
}
template <typename scalar_t>
__global__ void respond_kernel(int64_t *__restrict__ cluster, int64_t *proposal,
int64_t *__restrict row,
int64_t *__restrict__ col,
scalar_t *__restrict__ weight, size_t numel) {
const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
const size_t stride = blockDim.x * gridDim.x;
for (int64_t u = index; u < numel; u += stride) {
if (cluster[u] != -2)
continue; // Only vist red nodes.
bool has_unmatched_neighbor = false;
int64_t v_max = -1;
scalar_t w_max = 0;
for (int64_t i = row[u]; i < row[u + 1]; i++) {
auto v = col[i];
if (cluster[v] < 0)
has_unmatched_neighbor = true; // Unmatched neighbor found.
if (cluster[v] == -1 && proposal[v] == u && weight[i] >= w_max) {
// Find maximum weighted blue neighbhor v which proposed to u.
v_max = v;
w_max = weight[i];
}
}
if (v_max >= 0) {
cluster[u] = min(u, v_max); // Match neighbors.
cluster[v_max] = min(u, v_max);
}
if (!has_unmatched_neighbor)
cluster[u] = u;
}
}
void respond(at::Tensor cluster, at::Tensor proposal, at::Tensor row,
at::Tensor col, at::Tensor weight) {
AT_DISPATCH_ALL_TYPES(weight.scalar_type(), "respond_kernel", [&] {
respond_kernel<scalar_t><<<BLOCKS(cluster.numel()), THREADS>>>(
cluster.DATA_PTR<int64_t>(), proposal.DATA_PTR<int64_t>(),
row.DATA_PTR<int64_t>(), col.DATA_PTR<int64_t>(),
weight.DATA_PTR<scalar_t>(), cluster.numel());
});
}
cuda/rw.cpp deleted 100644 → 0
View file @ 0194ebb6
#include <torch/extension.h>
#define CHECK_CUDA(x) \
AT_ASSERTM(x.device().is_cuda(), #x " must be CUDA tensor")
#define IS_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " is not contiguous");
at::Tensor rw_cuda(at::Tensor row, at::Tensor col, at::Tensor start,
size_t walk_length, float p, float q, size_t num_nodes);
at::Tensor rw(at::Tensor row, at::Tensor col, at::Tensor start,
size_t walk_length, float p, float q, size_t num_nodes) {
CHECK_CUDA(row);
CHECK_CUDA(col);
CHECK_CUDA(start);
return rw_cuda(row, col, start, walk_length, p, q, num_nodes);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rw", &rw, "Random Walk Sampling (CUDA)");
}
cuda/rw_kernel.cu deleted 100644 → 0
View file @ 0194ebb6
#include <ATen/ATen.h>
#include "compat.cuh"
#include "utils.cuh"
#define THREADS 1024
#define BLOCKS(N) (N + THREADS - 1) / THREADS
__global__ void uniform_rw_kernel(
const int64_t *__restrict__ row, const int64_t *__restrict__ col,
const int64_t *__restrict__ deg, const int64_t *__restrict__ start,
const float *__restrict__ rand, int64_t *__restrict__ out,
const size_t walk_length, const size_t numel) {
const size_t index = blockIdx.x * blockDim.x + threadIdx.x;
const size_t stride = blockDim.x * gridDim.x;
for (ptrdiff_t n = index; n < numel; n += stride) {
out[n] = start[n];
for (ptrdiff_t l = 1; l <= walk_length; l++) {
auto i = (l - 1) * numel + n;
auto cur = out[i];
out[l * numel + n] = col[row[cur] + int64_t(rand[i] * deg[cur])];
}
}
}
at::Tensor rw_cuda(at::Tensor row, at::Tensor col, at::Tensor start,
size_t walk_length, float p, float q, size_t num_nodes) {
cudaSetDevice(row.get_device());
auto deg = degree(row, num_nodes);
row = at::cat({at::zeros(1, deg.options()), deg.cumsum(0)}, 0);
auto rand = at::rand({(int64_t)walk_length, start.size(0)},
start.options().dtype(at::kFloat));
auto out =
at::full({(int64_t)walk_length + 1, start.size(0)}, -1, start.options());
uniform_rw_kernel<<<BLOCKS(start.numel()), THREADS>>>(
row.DATA_PTR<int64_t>(), col.DATA_PTR<int64_t>(), deg.DATA_PTR<int64_t>(),
start.DATA_PTR<int64_t>(), rand.DATA_PTR<float>(),
out.DATA_PTR<int64_t>(), walk_length, start.numel());
return out.t().contiguous();
}
cuda/utils.cuh deleted 100644 → 0
View file @ 0194ebb6
#pragma once
#include <ATen/ATen.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::empty(row.size(0), row.options());
at::randperm_out(perm, row.size(0));
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);
}
script/conda.sh 0 → 100755
View file @ 80b99adb
#!/bin/bash
if [ "${TRAVIS_OS_NAME}" = "linux" ]; then
wget -nv https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
chmod +x miniconda.sh
./miniconda.sh -b
PATH=/home/travis/miniconda3/bin:${PATH}
fi
if [ "${TRAVIS_OS_NAME}" = "osx" ]; then
wget -nv https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh -O miniconda.sh
chmod +x miniconda.sh
./miniconda.sh -b
PATH=/Users/travis/miniconda3/bin:${PATH}
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ]; then
choco install openssl.light
choco install miniconda3
PATH=/c/tools/miniconda3/Scripts:$PATH
fi
conda update --yes conda
conda create --yes -n test python="${PYTHON_VERSION}"
script/cuda.sh 0 → 100755
View file @ 80b99adb
#!/bin/bash
if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "$IDX" = "cpu" ]; then
export TOOLKIT=cpuonly
fi
if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "$IDX" = "cu92" ]; then
export CUDA_SHORT=9.2
export CUDA=9.2.148-1
export UBUNTU_VERSION=ubuntu1604
export CUBLAS=cuda-cublas-dev-9-2
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "$IDX" = "cu100" ]; then
export CUDA_SHORT=10.0
export CUDA=10.0.130-1
export UBUNTU_VERSION=ubuntu1804
export CUBLAS=cuda-cublas-dev-10-0
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "$IDX" = "cu101" ]; then
export IDX=cu101
export CUDA_SHORT=10.1
export CUDA=10.1.105-1
export UBUNTU_VERSION=ubuntu1804
export CUBLAS=libcublas-dev
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "$IDX" = "cpu" ]; then
export TOOLKIT=cpuonly
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "$IDX" = "cu92" ]; then
export CUDA_SHORT=9.2
export CUDA_URL=https://developer.nvidia.com/compute/cuda/${CUDA_SHORT}/Prod2/local_installers2
export CUDA_FILE=cuda_${CUDA_SHORT}.148_win10
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "$IDX" = "cu100" ]; then
export CUDA_SHORT=10.0
export CUDA_URL=https://developer.nvidia.com/compute/cuda/${CUDA_SHORT}/Prod/local_installers
export CUDA_FILE=cuda_${CUDA_SHORT}.130_411.31_win10
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "$IDX" = "cu101" ]; then
export CUDA_SHORT=10.1
export CUDA_URL=https://developer.nvidia.com/compute/cuda/${CUDA_SHORT}/Prod/local_installers
export CUDA_FILE=cuda_${CUDA_SHORT}.105_418.96_win10.exe
export TOOLKIT="cudatoolkit=${CUDA_SHORT}"
fi
if [ "${TRAVIS_OS_NAME}" = "osx" ] && [ "$IDX" = "cpu" ]; then
export TOOLKIT=""
fi
if [ "${IDX}" = "cpu" ]; then
export FORCE_CPU=1
else
export FORCE_CUDA=1
fi
if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "${IDX}" != "cpu" ]; then
INSTALLER=cuda-repo-${UBUNTU_VERSION}_${CUDA}_amd64.deb
wget -nv "http://developer.download.nvidia.com/compute/cuda/repos/${UBUNTU_VERSION}/x86_64/${INSTALLER}"
sudo dpkg -i "${INSTALLER}"
wget -nv "https://developer.download.nvidia.com/compute/cuda/repos/${UBUNTU_VERSION}/x86_64/7fa2af80.pub"
sudo apt-key add 7fa2af80.pub
sudo apt update -qq
sudo apt install -y "cuda-core-${CUDA_SHORT/./-}" "cuda-cudart-dev-${CUDA_SHORT/./-}" "${CUBLAS}" "cuda-cusparse-dev-${CUDA_SHORT/./-}"
sudo apt clean
CUDA_HOME=/usr/local/cuda-${CUDA_SHORT}
LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
PATH=${CUDA_HOME}/bin:${PATH}
nvcc --version
fi
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${IDX}" != "cpu" ]; then
wget -nv "${CUDA_URL}/${CUDA_FILE}"
PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT}\" -Wait -NoNewWindow"
CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v${CUDA_SHORT}
PATH=${CUDA_HOME}/bin:$PATH
PATH=/c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/MSBuild/15.0/Bin:$PATH
nvcc --version
fi
# Fix Cuda9.2 on Windows: https://github.com/pytorch/pytorch/issues/6109
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${IDX}" = "cu92" ]; then
sed -i.bak -e '129,141d' "${CUDA_HOME}/include/crt/host_config.h"
fi
script/rename_wheel.py 0 → 100644
View file @ 80b99adb
import sys
import os
import os.path as osp
import glob
import shutil
idx = sys.argv[1]
assert idx in ['cpu', 'cu92', 'cu100', 'cu101']
dist_dir = osp.join(osp.dirname(osp.abspath(__file__)), '..', 'dist')
wheels = glob.glob(osp.join('dist', '**', '*.whl'), recursive=True)
for wheel in wheels:
if idx in wheel:
continue
paths = wheel.split(osp.sep)
names = paths[-1].split('-')
name = '-'.join(names[:-4] + ['latest+' + idx] + names[-3:])
shutil.copyfile(wheel, osp.join(*paths[:-1], name))
name = '-'.join(names[:-4] + [names[-4] + '+' + idx] + names[-3:])
os.rename(wheel, osp.join(*paths[:-1], name))
script/torch.sh 0 → 100755
View file @ 80b99adb
#!/bin/bash
# Fix "member may not be initialized" error on Windows: https://github.com/pytorch/pytorch/issues/27958
if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${IDX}" != "cpu" ]; then
sed -i.bak -e 's/constexpr/const/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/script/module.h
sed -i.bak -e 's/constexpr/const/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/torch/csrc/jit/argument_spec.h
sed -i.bak -e 's/return \*(this->value)/return \*((type\*)this->value)/g' /c/tools/miniconda3/envs/test/lib/site-packages/torch/include/pybind11/cast.h
fi
setup.py
View file @ 80b99adb
import os
import os.path as osp
import glob
from setuptools import setup, find_packages from setuptools import setup, find_packages
from sys import argv
import torch import torch
from torch.utils.cpp_extension import BuildExtension
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
TORCH_MAJOR = int(torch.__version__.split('.')[0]) WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
TORCH_MINOR = int(torch.__version__.split('.')[1]) if os.getenv('FORCE_CUDA', '0') == '1':
WITH_CUDA = True
if os.getenv('FORCE_CPU', '0') == '1':
WITH_CUDA = False
BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
def get_extensions():
Extension = CppExtension
define_macros = []
extra_compile_args = {'cxx': []}
if WITH_CUDA:
Extension = CUDAExtension
define_macros += [('WITH_CUDA', None)]
nvcc_flags = os.getenv('NVCC_FLAGS', '')
nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
nvcc_flags += ['-arch=sm_35', '--expt-relaxed-constexpr']
extra_compile_args['nvcc'] = nvcc_flags
extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
extensions = []
for main in main_files:
name = main.split(os.sep)[-1][:-4]
extra_compile_args = [] sources = [main]
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
extra_compile_args += ['-DVERSION_GE_1_3']
ext_modules = [ path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
CppExtension('torch_cluster.graclus_cpu', ['cpu/graclus.cpp'], if osp.exists(path):
extra_compile_args=extra_compile_args), sources += [path]
CppExtension('torch_cluster.grid_cpu', ['cpu/grid.cpp']),
CppExtension('torch_cluster.fps_cpu', ['cpu/fps.cpp'],
extra_compile_args=extra_compile_args),
CppExtension('torch_cluster.rw_cpu', ['cpu/rw.cpp'],
extra_compile_args=extra_compile_args),
CppExtension('torch_cluster.sampler_cpu', ['cpu/sampler.cpp'],
extra_compile_args=extra_compile_args),
]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
GPU = True path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
for arg in argv: if WITH_CUDA and osp.exists(path):
if arg == '--cpu': sources += [path]
GPU = False
argv.remove(arg)
if CUDA_HOME is not None and GPU: extension = Extension(
ext_modules += [ 'torch_cluster._' + name,
CUDAExtension('torch_cluster.graclus_cuda', sources,
['cuda/graclus.cpp', 'cuda/graclus_kernel.cu'], include_dirs=[extensions_dir],
extra_compile_args=extra_compile_args), define_macros=define_macros,
CUDAExtension('torch_cluster.grid_cuda', extra_compile_args=extra_compile_args,
['cuda/grid.cpp', 'cuda/grid_kernel.cu'], )
extra_compile_args=extra_compile_args), extensions += [extension]
CUDAExtension('torch_cluster.fps_cuda',
['cuda/fps.cpp', 'cuda/fps_kernel.cu'], return extensions
extra_compile_args=extra_compile_args),
CUDAExtension('torch_cluster.nearest_cuda',
['cuda/nearest.cpp', 'cuda/nearest_kernel.cu'],
extra_compile_args=extra_compile_args),
CUDAExtension('torch_cluster.knn_cuda',
['cuda/knn.cpp', 'cuda/knn_kernel.cu'],
extra_compile_args=extra_compile_args),
CUDAExtension('torch_cluster.radius_cuda',
['cuda/radius.cpp', 'cuda/radius_kernel.cu'],
extra_compile_args=extra_compile_args),
CUDAExtension('torch_cluster.rw_cuda',
['cuda/rw.cpp', 'cuda/rw_kernel.cu'],
extra_compile_args=extra_compile_args),
]
__version__ = '1.4.5'
url = 'https://github.com/rusty1s/pytorch_cluster'
install_requires = ['scipy'] install_requires = ['scipy']
setup_requires = ['pytest-runner'] setup_requires = ['pytest-runner']
...@@ -63,23 +63,26 @@ tests_require = ['pytest', 'pytest-cov'] ...@@ -63,23 +63,26 @@ tests_require = ['pytest', 'pytest-cov']
setup( setup(
name='torch_cluster', name='torch_cluster',
version=__version__, version='1.5.0',
description=('PyTorch Extension Library of Optimized Graph Cluster '
'Algorithms'),
author='Matthias Fey', author='Matthias Fey',
author_email='matthias.fey@tu-dortmund.de', author_email='matthias.fey@tu-dortmund.de',
url=url, url='https://github.com/rusty1s/pytorch_cluster',
download_url='{}/archive/{}.tar.gz'.format(url, __version__), description=('PyTorch Extension Library of Optimized Graph Cluster '
'Algorithms'),
keywords=[ keywords=[
'pytorch', 'pytorch',
'geometric-deep-learning', 'geometric-deep-learning',
'graph-neural-networks', 'graph-neural-networks',
'cluster-algorithms', 'cluster-algorithms',
], ],
license='MIT',
python_requires='>=3.6',
install_requires=install_requires, install_requires=install_requires,
setup_requires=setup_requires, setup_requires=setup_requires,
tests_require=tests_require, tests_require=tests_require,
ext_modules=ext_modules, ext_modules=get_extensions() if not BUILD_DOCS else [],
cmdclass=cmdclass, cmdclass={
'build_ext': BuildExtension.with_options(no_python_abi_suffix=True)
},
packages=find_packages(), packages=find_packages(),
) )
test/utils.py
View file @ 80b99adb
import torch import torch
from torch.testing import get_all_dtypes
dtypes = get_all_dtypes()
dtypes.remove(torch.half)
dtypes.remove(torch.bool)
if hasattr(torch, 'bfloat16'):
dtypes.remove(torch.bfloat16)
dtypes = [torch.float, torch.double, torch.int, torch.long]
grad_dtypes = [torch.float, torch.double] grad_dtypes = [torch.float, torch.double]
devices = [torch.device('cpu')] devices = [torch.device('cpu')]
if torch.cuda.is_available(): if torch.cuda.is_available():
devices += [torch.device('cuda:{}'.format(torch.cuda.current_device()))] devices += [torch.device(f'cuda:{torch.cuda.current_device()}')]
def tensor(x, dtype, device): def tensor(x, dtype, device):
......
torch_cluster/__init__.py
View file @ 80b99adb
from .graclus import graclus_cluster import importlib
from .grid import grid_cluster import os.path as osp
from .fps import fps
from .nearest import nearest import torch
from .knn import knn, knn_graph
from .radius import radius, radius_graph __version__ = '1.5.0'
from .rw import random_walk expected_torch_version = (1, 4)
from .sampler import neighbor_sampler
try:
__version__ = '1.4.5' for library in [
'_version', '_grid', '_graclus', '_fps', '_rw', '_sampler',
'_nearest', '_knn', '_radius'
]:
torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
library, [osp.dirname(__file__)]).origin)
except OSError as e:
major, minor = [int(x) for x in torch.__version__.split('.')[:2]]
t_major, t_minor = expected_torch_version
if major != t_major or (major == t_major and minor != t_minor):
raise RuntimeError(
f'Expected PyTorch version {t_major}.{t_minor} but found '
f'version {major}.{minor}.')
raise OSError(e)
if torch.version.cuda is not None: # pragma: no cover
cuda_version = torch.ops.torch_cluster.cuda_version()
if cuda_version == -1:
major = minor = 0
elif cuda_version < 10000:
major, minor = int(str(cuda_version)[0]), int(str(cuda_version)[2])
else:
major, minor = int(str(cuda_version)[0:2]), int(str(cuda_version)[3])
t_major, t_minor = [int(x) for x in torch.version.cuda.split('.')]
if t_major != major or t_minor != minor:
raise RuntimeError(
f'Detected that PyTorch and torch_cluster were compiled with '
f'different CUDA versions. PyTorch has CUDA version '
f'{t_major}.{t_minor} and torch_cluster has CUDA version '
f'{major}.{minor}. Please reinstall the torch_cluster that '
f'matches your PyTorch install.')
from .graclus import graclus_cluster # noqa
from .grid import grid_cluster # noqa
from .fps import fps # noqa
from .nearest import nearest # noqa
from .knn import knn, knn_graph # noqa
from .radius import radius, radius_graph # noqa
from .rw import random_walk # noqa
from .sampler import neighbor_sampler # noqa
__all__ = [ __all__ = [
'graclus_cluster', 'graclus_cluster',
......
torch_cluster/fps.py
View file @ 80b99adb
import torch from typing import Optional
import torch_cluster.fps_cpu
if torch.cuda.is_available(): import torch
import torch_cluster.fps_cuda
def fps(x, batch=None, ratio=0.5, random_start=True): @torch.jit.script
def fps(src: torch.Tensor, batch: Optional[torch.Tensor] = None,
ratio: float = 0.5, random_start: bool = True) -> torch.Tensor:
r""""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature r""""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature
Learning on Point Sets in a Metric Space" Learning on Point Sets in a Metric Space"
<https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the <https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the
most distant point with regard to the rest points. most distant point with regard to the rest points.
Args: Args:
x (Tensor): Node feature matrix src (Tensor): Point feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`. :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
batch (LongTensor, optional): Batch vector batch (LongTensor, optional): Batch vector
:math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
...@@ -23,28 +23,26 @@ def fps(x, batch=None, ratio=0.5, random_start=True): ...@@ -23,28 +23,26 @@ def fps(x, batch=None, ratio=0.5, random_start=True):
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
.. testsetup:: .. code-block:: python
import torch import torch
from torch_cluster import fps from torch_cluster import fps
.. testcode:: src = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch = torch.tensor([0, 0, 0, 0])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) index = fps(src, batch, ratio=0.5)
>>> batch = torch.tensor([0, 0, 0, 0])
>>> index = fps(x, batch, ratio=0.5)
""" """
if batch is None: if batch is not None:
batch = x.new_zeros(x.size(0), dtype=torch.long) assert src.size(0) == batch.numel()
batch_size = int(batch.max()) + 1
x = x.view(-1, 1) if x.dim() == 1 else x deg = src.new_zeros(batch_size, dtype=torch.long)
deg.scatter_add_(0, batch, torch.ones_like(batch))
assert x.dim() == 2 and batch.dim() == 1 ptr = deg.new_zeros(batch_size + 1)
assert x.size(0) == batch.size(0) torch.cumsum(deg, 0, out=ptr[1:])
assert ratio > 0 and ratio < 1
if x.is_cuda:
return torch_cluster.fps_cuda.fps(x, batch, ratio, random_start)
else: else:
return torch_cluster.fps_cpu.fps(x, batch, ratio, random_start) ptr = torch.tensor([0, src.size(0)], device=src.device)
return torch.ops.torch_cluster.fps(src, ptr, ratio, random_start)
torch_cluster/graclus.py
View file @ 80b99adb
import torch from typing import Optional
import torch_cluster.graclus_cpu
if torch.cuda.is_available(): import torch
import torch_cluster.graclus_cuda
def graclus_cluster(row, col, weight=None, num_nodes=None): @torch.jit.script
def graclus_cluster(row: torch.Tensor, col: torch.Tensor,
weight: Optional[torch.Tensor] = None,
num_nodes: Optional[int] = None) -> torch.Tensor:
"""A greedy clustering algorithm of picking an unmarked vertex and matching """A greedy clustering algorithm of picking an unmarked vertex and matching
it with one its unmarked neighbors (that maximizes its edge weight). it with one its unmarked neighbors (that maximizes its edge weight).
...@@ -17,25 +18,42 @@ def graclus_cluster(row, col, weight=None, num_nodes=None): ...@@ -17,25 +18,42 @@ def graclus_cluster(row, col, weight=None, num_nodes=None):
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
Examples:: .. code-block:: python
>>> row = torch.tensor([0, 1, 1, 2]) import torch
>>> col = torch.tensor([1, 0, 2, 1]) from torch_cluster import graclus_cluster
>>> weight = torch.Tensor([1, 1, 1, 1])
>>> cluster = graclus_cluster(row, col, weight) row = torch.tensor([0, 1, 1, 2])
col = torch.tensor([1, 0, 2, 1])
weight = torch.Tensor([1, 1, 1, 1])
cluster = graclus_cluster(row, col, weight)
""" """
if num_nodes is None: if num_nodes is None:
num_nodes = max(row.max().item(), col.max().item()) + 1 num_nodes = max(int(row.max()), int(col.max())) + 1
# Remove self-loops.
mask = row == col
row, col = row[mask], col[mask]
if row.is_cuda: if weight is not None:
op = torch_cluster.graclus_cuda weight = weight[mask]
else:
op = torch_cluster.graclus_cpu
# Randomly shuffle nodes.
if weight is None: if weight is None:
cluster = op.graclus(row, col, num_nodes) perm = torch.randperm(row.size(0), dtype=torch.long, device=row.device)
else: row, col = row[perm], col[perm]
cluster = op.weighted_graclus(row, col, weight, num_nodes)
# To CSR.
perm = torch.argsort(row)
row, col = row[perm], col[perm]
if weight is not None:
weight = weight[perm]
deg = row.new_zeros(num_nodes)
deg.scatter_add_(0, row, torch.ones_like(row))
rowptr = row.new_zeros(num_nodes + 1)
torch.cumsum(deg, 0, out=rowptr[1:])
return cluster return torch.ops.torch_cluster.graclus(rowptr, col, weight)
torch_cluster/grid.py
View file @ 80b99adb
import torch from typing import Optional
import torch_cluster.grid_cpu
if torch.cuda.is_available(): import torch
import torch_cluster.grid_cuda
def grid_cluster(pos, size, start=None, end=None): @torch.jit.script
def grid_cluster(pos: torch.Tensor, size: torch.Tensor,
start: Optional[torch.Tensor] = None,
end: Optional[torch.Tensor] = None) -> torch.Tensor:
"""A clustering algorithm, which overlays a regular grid of user-defined """A clustering algorithm, which overlays a regular grid of user-defined
size over a point cloud and clusters all points within a voxel. size over a point cloud and clusters all points within a voxel.
...@@ -19,22 +20,13 @@ def grid_cluster(pos, size, start=None, end=None): ...@@ -19,22 +20,13 @@ def grid_cluster(pos, size, start=None, end=None):
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
Examples:: .. code-block:: python
>>> pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
>>> size = torch.Tensor([5, 5])
>>> cluster = grid_cluster(pos, size)
"""
pos = pos.unsqueeze(-1) if pos.dim() == 1 else pos import torch
start = pos.t().min(dim=1)[0] if start is None else start from torch_cluster import grid_cluster
end = pos.t().max(dim=1)[0] if end is None else end
if pos.is_cuda: pos = torch.Tensor([[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]])
op = torch_cluster.grid_cuda size = torch.Tensor([5, 5])
else: cluster = grid_cluster(pos, size)
op = torch_cluster.grid_cpu """
return torch.ops.torch_cluster.grid(pos, size, start, end)
cluster = op.grid(pos, size, start, end)
return cluster
torch_cluster/knn.py
View file @ 80b99adb
from typing import Optional
import torch import torch
import scipy.spatial import scipy.spatial
if torch.cuda.is_available():
import torch_cluster.knn_cuda
def knn(x, y, k, batch_x=None, batch_y=None, cosine=False): def knn(x: torch.Tensor, y: torch.Tensor, k: int,
batch_x: Optional[torch.Tensor] = None,
batch_y: Optional[torch.Tensor] = None,
cosine: bool = False) -> torch.Tensor:
r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
:obj:`x`. :obj:`x`.
...@@ -27,66 +29,91 @@ def knn(x, y, k, batch_x=None, batch_y=None, cosine=False): ...@@ -27,66 +29,91 @@ def knn(x, y, k, batch_x=None, batch_y=None, cosine=False):
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
.. testsetup:: .. code-block:: python
import torch import torch
from torch_cluster import knn from torch_cluster import knn
.. testcode:: x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.tensor([0, 0, 0, 0])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) y = torch.Tensor([[-1, 0], [1, 0]])
>>> batch_x = torch.tensor([0, 0, 0, 0]) batch_x = torch.tensor([0, 0])
>>> y = torch.Tensor([[-1, 0], [1, 0]]) assign_index = knn(x, y, 2, batch_x, batch_y)
>>> batch_x = torch.tensor([0, 0])
>>> assign_index = knn(x, y, 2, batch_x, batch_y)
""" """
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
x = x.view(-1, 1) if x.dim() == 1 else x x = x.view(-1, 1) if x.dim() == 1 else x
y = y.view(-1, 1) if y.dim() == 1 else y y = y.view(-1, 1) if y.dim() == 1 else y
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
if x.is_cuda: if x.is_cuda:
return torch_cluster.knn_cuda.knn(x, y, k, batch_x, batch_y, cosine) if batch_x is not None:
assert x.size(0) == batch_x.numel()
if cosine: batch_size = int(batch_x.max()) + 1
raise NotImplementedError('Cosine distance not implemented for CPU')
deg = x.new_zeros(batch_size, dtype=torch.long)
# Rescale x and y. deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))
min_xy = min(x.min().item(), y.min().item())
x, y = x - min_xy, y - min_xy ptr_x = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_x[1:])
max_xy = max(x.max().item(), y.max().item()) else:
x, y, = x / max_xy, y / max_xy ptr_x = torch.tensor([0, x.size(0)], device=x.device)
# Concat batch/features to ensure no cross-links between examples exist. if batch_y is not None:
x = torch.cat([x, 2 * x.size(1) * batch_x.view(-1, 1).to(x.dtype)], dim=-1) assert y.size(0) == batch_y.numel()
y = torch.cat([y, 2 * y.size(1) * batch_y.view(-1, 1).to(y.dtype)], dim=-1) batch_size = int(batch_y.max()) + 1
tree = scipy.spatial.cKDTree(x.detach().numpy()) deg = y.new_zeros(batch_size, dtype=torch.long)
dist, col = tree.query(y.detach().cpu(), k=k, deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))
distance_upper_bound=x.size(1))
dist = torch.from_numpy(dist).to(x.dtype) ptr_y = deg.new_zeros(batch_size + 1)
col = torch.from_numpy(col).to(torch.long) torch.cumsum(deg, 0, out=ptr_y[1:])
row = torch.arange(col.size(0), dtype=torch.long).view(-1, 1).repeat(1, k) else:
mask = ~torch.isinf(dist).view(-1) ptr_y = torch.tensor([0, y.size(0)], device=y.device)
row, col = row.view(-1)[mask], col.view(-1)[mask]
return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine)
return torch.stack([row, col], dim=0) else:
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
def knn_graph(x, k, batch=None, loop=False, flow='source_to_target',
cosine=False): if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
if cosine:
raise NotImplementedError('`cosine` argument not supported on CPU')
# Translate and rescale x and y to [0, 1].
min_xy = min(x.min().item(), y.min().item())
x, y = x - min_xy, y - min_xy
max_xy = max(x.max().item(), y.max().item())
x.div_(max_xy)
y.div_(max_xy)
# Concat batch/features to ensure no cross-links between examples.
x = torch.cat([x, 2 * x.size(1) * batch_x.view(-1, 1).to(x.dtype)], -1)
y = torch.cat([y, 2 * y.size(1) * batch_y.view(-1, 1).to(y.dtype)], -1)
tree = scipy.spatial.cKDTree(x.detach().numpy())
dist, col = tree.query(y.detach().cpu(), k=k,
distance_upper_bound=x.size(1))
dist = torch.from_numpy(dist).to(x.dtype)
col = torch.from_numpy(col).to(torch.long)
row = torch.arange(col.size(0), dtype=torch.long)
row = row.view(-1, 1).repeat(1, k)
mask = ~torch.isinf(dist).view(-1)
row, col = row.view(-1)[mask], col.view(-1)[mask]
return torch.stack([row, col], dim=0)
def knn_graph(x: torch.Tensor, k: int, batch: Optional[torch.Tensor] = None,
loop: bool = False, flow: str = 'source_to_target',
cosine: bool = False) -> torch.Tensor:
r"""Computes graph edges to the nearest :obj:`k` points. r"""Computes graph edges to the nearest :obj:`k` points.
Args: Args:
...@@ -107,16 +134,14 @@ def knn_graph(x, k, batch=None, loop=False, flow='source_to_target', ...@@ -107,16 +134,14 @@ def knn_graph(x, k, batch=None, loop=False, flow='source_to_target',
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
.. testsetup:: .. code-block:: python
import torch import torch
from torch_cluster import knn_graph from torch_cluster import knn_graph
.. testcode:: x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch = torch.tensor([0, 0, 0, 0])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) edge_index = knn_graph(x, k=2, batch=batch, loop=False)
>>> batch = torch.tensor([0, 0, 0, 0])
>>> edge_index = knn_graph(x, k=2, batch=batch, loop=False)
""" """
assert flow in ['source_to_target', 'target_to_source'] assert flow in ['source_to_target', 'target_to_source']
......
torch_cluster/nearest.py
View file @ 80b99adb
from typing import Optional
import torch import torch
import scipy.cluster import scipy.cluster
if torch.cuda.is_available():
import torch_cluster.nearest_cuda
def nearest(x, y, batch_x=None, batch_y=None): def nearest(x: torch.Tensor, y: torch.Tensor,
batch_x: Optional[torch.Tensor] = None,
batch_y: Optional[torch.Tensor] = None) -> torch.Tensor:
r"""Clusters points in :obj:`x` together which are nearest to a given query r"""Clusters points in :obj:`x` together which are nearest to a given query
point in :obj:`y`. point in :obj:`y`.
...@@ -21,49 +22,74 @@ def nearest(x, y, batch_x=None, batch_y=None): ...@@ -21,49 +22,74 @@ def nearest(x, y, batch_x=None, batch_y=None):
:math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
node to a specific example. (default: :obj:`None`) node to a specific example. (default: :obj:`None`)
.. testsetup:: :rtype: :class:`LongTensor`
.. code-block:: python
import torch import torch
from torch_cluster import nearest from torch_cluster import nearest
.. testcode:: x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.tensor([0, 0, 0, 0])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) y = torch.Tensor([[-1, 0], [1, 0]])
>>> batch_x = torch.tensor([0, 0, 0, 0]) batch_y = torch.tensor([0, 0])
>>> y = torch.Tensor([[-1, 0], [1, 0]]) cluster = nearest(x, y, batch_x, batch_y)
>>> batch_y = torch.tensor([0, 0])
>>> cluster = nearest(x, y, batch_x, batch_y)
""" """
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
x = x.view(-1, 1) if x.dim() == 1 else x x = x.view(-1, 1) if x.dim() == 1 else x
y = y.view(-1, 1) if y.dim() == 1 else y y = y.view(-1, 1) if y.dim() == 1 else y
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
if x.is_cuda: if x.is_cuda:
return torch_cluster.nearest_cuda.nearest(x, y, batch_x, batch_y) if batch_x is not None:
assert x.size(0) == batch_x.numel()
# Rescale x and y. batch_size = int(batch_x.max()) + 1
min_xy = min(x.min().item(), y.min().item())
x, y = x - min_xy, y - min_xy deg = x.new_zeros(batch_size, dtype=torch.long)
deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))
max_xy = max(x.max().item(), y.max().item())
x, y, = x / max_xy, y / max_xy ptr_x = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_x[1:])
# Concat batch/features to ensure no cross-links between examples exist. else:
x = torch.cat([x, 2 * x.size(1) * batch_x.view(-1, 1).to(x.dtype)], dim=-1) ptr_x = torch.tensor([0, x.size(0)], device=x.device)
y = torch.cat([y, 2 * y.size(1) * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
if batch_y is not None:
return torch.from_numpy( assert y.size(0) == batch_y.numel()
scipy.cluster.vq.vq(x.detach().cpu(), batch_size = int(batch_y.max()) + 1
y.detach().cpu())[0]).to(torch.long)
deg = y.new_zeros(batch_size, dtype=torch.long)
deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))
ptr_y = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_y[1:])
else:
ptr_y = torch.tensor([0, y.size(0)], device=y.device)
return torch.ops.torch_cluster.nearest(x, y, ptr_x, ptr_y)
else:
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
# Translate and rescale x and y to [0, 1].
min_xy = min(x.min().item(), y.min().item())
x, y = x - min_xy, y - min_xy
max_xy = max(x.max().item(), y.max().item())
x.div_(max_xy)
y.div_(max_xy)
# Concat batch/features to ensure no cross-links between examples.
x = torch.cat([x, 2 * x.size(1) * batch_x.view(-1, 1).to(x.dtype)], -1)
y = torch.cat([y, 2 * y.size(1) * batch_y.view(-1, 1).to(y.dtype)], -1)
return torch.from_numpy(
scipy.cluster.vq.vq(x.detach().cpu(),
y.detach().cpu())[0]).to(torch.long)
torch_cluster/radius.py
View file @ 80b99adb
from typing import Optional
import torch import torch
import scipy.spatial import scipy.spatial
if torch.cuda.is_available():
import torch_cluster.radius_cuda
def sample(col, count): @torch.jit.script
def sample(col: torch.Tensor, count: int) -> torch.Tensor:
if col.size(0) > count: if col.size(0) > count:
col = col[torch.randperm(col.size(0))][:count] col = col[torch.randperm(col.size(0))][:count]
return col return col
def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32): def radius(x: torch.Tensor, y: torch.Tensor, r: float,
batch_x: Optional[torch.Tensor] = None,
batch_y: Optional[torch.Tensor] = None,
max_num_neighbors: int = 32) -> torch.Tensor:
r"""Finds for each element in :obj:`y` all points in :obj:`x` within r"""Finds for each element in :obj:`y` all points in :obj:`x` within
distance :obj:`r`. distance :obj:`r`.
...@@ -30,56 +33,76 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32): ...@@ -30,56 +33,76 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
max_num_neighbors (int, optional): The maximum number of neighbors to max_num_neighbors (int, optional): The maximum number of neighbors to
return for each element in :obj:`y`. (default: :obj:`32`) return for each element in :obj:`y`. (default: :obj:`32`)
:rtype: :class:`LongTensor` .. code-block:: python
.. testsetup::
import torch import torch
from torch_cluster import radius from torch_cluster import radius
.. testcode:: x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.Tensor([[-1, 0], [1, 0]])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) batch_y = torch.tensor([0, 0])
>>> batch_x = torch.tensor([0, 0, 0, 0]) assign_index = radius(x, y, 1.5, batch_x, batch_y)
>>> y = torch.Tensor([[-1, 0], [1, 0]])
>>> batch_y = torch.tensor([0, 0])
>>> assign_index = radius(x, y, 1.5, batch_x, batch_y)
""" """
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
x = x.view(-1, 1) if x.dim() == 1 else x x = x.view(-1, 1) if x.dim() == 1 else x
y = y.view(-1, 1) if y.dim() == 1 else y y = y.view(-1, 1) if y.dim() == 1 else y
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
if x.is_cuda: if x.is_cuda:
return torch_cluster.radius_cuda.radius(x, y, r, batch_x, batch_y, if batch_x is not None:
max_num_neighbors) assert x.size(0) == batch_x.numel()
batch_size = int(batch_x.max()) + 1
x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1) deg = x.new_zeros(batch_size, dtype=torch.long)
deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))
tree = scipy.spatial.cKDTree(x.detach().numpy())
col = tree.query_ball_point(y.detach().numpy(), r) ptr_x = deg.new_zeros(batch_size + 1)
col = [sample(torch.tensor(c), max_num_neighbors) for c in col] torch.cumsum(deg, 0, out=ptr_x[1:])
row = [torch.full_like(c, i) for i, c in enumerate(col)] else:
row, col = torch.cat(row, dim=0), torch.cat(col, dim=0) ptr_x = torch.tensor([0, x.size(0)], device=x.device)
mask = col < int(tree.n)
return torch.stack([row[mask], col[mask]], dim=0) if batch_y is not None:
assert y.size(0) == batch_y.numel()
batch_size = int(batch_y.max()) + 1
def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32,
flow='source_to_target'): deg = y.new_zeros(batch_size, dtype=torch.long)
deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))
ptr_y = deg.new_zeros(batch_size + 1)
torch.cumsum(deg, 0, out=ptr_y[1:])
else:
ptr_y = torch.tensor([0, y.size(0)], device=y.device)
return torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
max_num_neighbors)
else:
if batch_x is None:
batch_x = x.new_zeros(x.size(0), dtype=torch.long)
if batch_y is None:
batch_y = y.new_zeros(y.size(0), dtype=torch.long)
assert x.dim() == 2 and batch_x.dim() == 1
assert y.dim() == 2 and batch_y.dim() == 1
assert x.size(1) == y.size(1)
assert x.size(0) == batch_x.size(0)
assert y.size(0) == batch_y.size(0)
x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)
tree = scipy.spatial.cKDTree(x.detach().numpy())
col = tree.query_ball_point(y.detach().numpy(), r)
col = [sample(torch.tensor(c), max_num_neighbors) for c in col]
row = [torch.full_like(c, i) for i, c in enumerate(col)]
row, col = torch.cat(row, dim=0), torch.cat(col, dim=0)
mask = col < int(tree.n)
return torch.stack([row[mask], col[mask]], dim=0)
def radius_graph(x: torch.Tensor, r: float,
batch: Optional[torch.Tensor] = None, loop: bool = False,
max_num_neighbors: int = 32,
flow: str = 'source_to_target') -> torch.Tensor:
r"""Computes graph edges to all points within a given distance. r"""Computes graph edges to all points within a given distance.
Args: Args:
...@@ -99,16 +122,14 @@ def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32, ...@@ -99,16 +122,14 @@ def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32,
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
.. testsetup:: .. code-block:: python
import torch import torch
from torch_cluster import radius_graph from torch_cluster import radius_graph
.. testcode:: x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch = torch.tensor([0, 0, 0, 0])
>>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) edge_index = radius_graph(x, r=1.5, batch=batch, loop=False)
>>> batch = torch.tensor([0, 0, 0, 0])
>>> edge_index = radius_graph(x, r=1.5, batch=batch, loop=False)
""" """
assert flow in ['source_to_target', 'target_to_source'] assert flow in ['source_to_target', 'target_to_source']
......
torch_cluster/rw.py
View file @ 80b99adb
import warnings import warnings
from typing import Optional
import torch import torch
import torch_cluster.rw_cpu
if torch.cuda.is_available():
import torch_cluster.rw_cuda
@torch.jit.script
def random_walk(row, col, start, walk_length, p=1, q=1, coalesced=False, def random_walk(row: torch.Tensor, col: torch.Tensor, start: torch.Tensor,
num_nodes=None): walk_length: int, p: float = 1, q: float = 1,
coalesced: bool = False, num_nodes: Optional[int] = None):
"""Samples random walks of length :obj:`walk_length` from all node indices """Samples random walks of length :obj:`walk_length` from all node indices
in :obj:`start` in the graph given by :obj:`(row, col)` as described in the in :obj:`start` in the graph given by :obj:`(row, col)` as described in the
`"node2vec: Scalable Feature Learning for Networks" `"node2vec: Scalable Feature Learning for Networks"
...@@ -33,22 +32,21 @@ def random_walk(row, col, start, walk_length, p=1, q=1, coalesced=False, ...@@ -33,22 +32,21 @@ def random_walk(row, col, start, walk_length, p=1, q=1, coalesced=False,
:rtype: :class:`LongTensor` :rtype: :class:`LongTensor`
""" """
if num_nodes is None: if num_nodes is None:
num_nodes = max(row.max(), col.max()).item() + 1 num_nodes = max(int(row.max()), int(col.max())) + 1
if coalesced: if coalesced:
_, perm = torch.sort(row * num_nodes + col) perm = torch.argsort(row * num_nodes + col)
row, col = row[perm], col[perm] row, col = row[perm], col[perm]
if p != 1 or q != 1: # pragma: no cover deg = row.new_zeros(num_nodes)
deg.scatter_add_(0, row, torch.ones_like(row))
rowptr = row.new_zeros(num_nodes + 1)
torch.cumsum(deg, 0, out=rowptr[1:])
if p != 1. or q != 1.: # pragma: no cover
warnings.warn('Parameters `p` and `q` are not supported yet and will' warnings.warn('Parameters `p` and `q` are not supported yet and will'
'be restored to their default values `p=1` and `q=1`.') 'be restored to their default values `p=1` and `q=1`.')
p = q = 1 p = q = 1.
start = start.flatten()
if row.is_cuda: # pragma: no cover return torch.ops.torch_cluster.random_walk(rowptr, col, start, walk_length,
return torch_cluster.rw_cuda.rw(row, col, start, walk_length, p, q, p, q)
num_nodes)
else:
return torch_cluster.rw_cpu.rw(row, col, start, walk_length, p, q,
num_nodes)
torch_cluster/sampler.py
View file @ 80b99adb
import torch_cluster.sampler_cpu import torch
def neighbor_sampler(start, cumdeg, size): @torch.jit.script
def neighbor_sampler(start: torch.Tensor, rowptr: torch.Tensor, size: float):
assert not start.is_cuda assert not start.is_cuda
factor = 1 factor: float = -1.
if isinstance(size, float): count: int = -1
if size <= 1:
factor = size factor = size
size = 2147483647 assert factor > 0
else:
count = int(size)
op = torch_cluster.sampler_cpu.neighbor_sampler return torch.ops.torch_cluster.neighbor_sampler(start, rowptr, count,
return op(start, cumdeg, size, factor) factor)
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