Skip to content

GitLab

Commit 676aa072 authored by rusty1s's avatar rusty1s
Browse files

pytorch 1.3 support

parent 0dd1186d
Show whitespace changes
Inline Side-by-side
Showing with 77 additions and 56 deletions
.travis.yml
View file @ 676aa072
......@@ -17,7 +17,7 @@ before_install:
- export CXX="g++-4.9"
install:
- pip install numpy
- pip install -q torch -f https://download.pytorch.org/whl/nightly/cpu/torch.html
- 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
......
cpu/basis.cpp
View file @ 676aa072
#include <torch/extension.h>
#include "compat.h"
template <typename scalar_t> inline scalar_t linear(scalar_t v, int64_t k_mod) {
return 1 - v - k_mod + 2 * v * k_mod;
}
......@@ -34,11 +36,11 @@ template <typename scalar_t> inline scalar_t cubic(scalar_t v, int64_t k_mod) {
\
AT_DISPATCH_FLOATING_TYPES( \
PSEUDO.scalar_type(), "basis_forward_##M", [&] { \
auto pseudo_data = PSEUDO.data<scalar_t>(); \
auto kernel_size_data = KERNEL_SIZE.data<int64_t>(); \
auto is_open_spline_data = IS_OPEN_SPLINE.data<uint8_t>(); \
auto basis_data = basis.data<scalar_t>(); \
auto weight_index_data = weight_index.data<int64_t>(); \
auto pseudo_data = PSEUDO.DATA_PTR<scalar_t>(); \
auto kernel_size_data = KERNEL_SIZE.DATA_PTR<int64_t>(); \
auto is_open_spline_data = IS_OPEN_SPLINE.DATA_PTR<uint8_t>(); \
auto basis_data = basis.DATA_PTR<scalar_t>(); \
auto weight_index_data = weight_index.DATA_PTR<int64_t>(); \
\
int64_t k, wi, wi_offset; \
scalar_t b; \
......@@ -126,11 +128,11 @@ inline scalar_t grad_cubic(scalar_t v, int64_t k_mod) {
\
AT_DISPATCH_FLOATING_TYPES( \
PSEUDO.scalar_type(), "basis_backward_##M", [&] { \
auto grad_basis_data = GRAD_BASIS.data<scalar_t>(); \
auto pseudo_data = PSEUDO.data<scalar_t>(); \
auto kernel_size_data = KERNEL_SIZE.data<int64_t>(); \
auto is_open_spline_data = IS_OPEN_SPLINE.data<uint8_t>(); \
auto grad_pseudo_data = grad_pseudo.data<scalar_t>(); \
auto grad_basis_data = GRAD_BASIS.DATA_PTR<scalar_t>(); \
auto pseudo_data = PSEUDO.DATA_PTR<scalar_t>(); \
auto kernel_size_data = KERNEL_SIZE.DATA_PTR<int64_t>(); \
auto is_open_spline_data = IS_OPEN_SPLINE.DATA_PTR<uint8_t>(); \
auto grad_pseudo_data = grad_pseudo.DATA_PTR<scalar_t>(); \
\
scalar_t g, tmp; \
\
......
cpu/compat.h 0 → 100644
View file @ 676aa072
#ifdef VERSION_GE_1_3
#define DATA_PTR data_ptr
#else
#define DATA_PTR data
#endif
cpu/weighting.cpp
View file @ 676aa072
#include <torch/extension.h>
#include "compat.h"
at::Tensor weighting_fw(at::Tensor x, at::Tensor weight, at::Tensor basis,
at::Tensor weight_index) {
auto E = x.size(0), M_in = x.size(1), M_out = weight.size(2);
......@@ -7,11 +9,11 @@ at::Tensor weighting_fw(at::Tensor x, at::Tensor weight, at::Tensor basis,
auto out = at::empty({E, M_out}, x.options());
AT_DISPATCH_FLOATING_TYPES(out.scalar_type(), "weighting_fw", [&] {
auto x_data = x.data<scalar_t>();
auto weight_data = weight.data<scalar_t>();
auto basis_data = basis.data<scalar_t>();
auto weight_index_data = weight_index.data<int64_t>();
auto out_data = out.data<scalar_t>();
auto x_data = x.DATA_PTR<scalar_t>();
auto weight_data = weight.DATA_PTR<scalar_t>();
auto basis_data = basis.DATA_PTR<scalar_t>();
auto weight_index_data = weight_index.DATA_PTR<int64_t>();
auto out_data = out.DATA_PTR<scalar_t>();
scalar_t v;
......@@ -44,11 +46,11 @@ at::Tensor weighting_bw_x(at::Tensor grad_out, at::Tensor weight,
auto grad_x = at::zeros({E, M_in}, grad_out.options());
AT_DISPATCH_FLOATING_TYPES(grad_out.scalar_type(), "weighting_bw_x", [&] {
auto grad_out_data = grad_out.data<scalar_t>();
auto weight_data = weight.data<scalar_t>();
auto basis_data = basis.data<scalar_t>();
auto weight_index_data = weight_index.data<int64_t>();
auto grad_x_data = grad_x.data<scalar_t>();
auto grad_out_data = grad_out.DATA_PTR<scalar_t>();
auto weight_data = weight.DATA_PTR<scalar_t>();
auto basis_data = basis.DATA_PTR<scalar_t>();
auto weight_index_data = weight_index.DATA_PTR<int64_t>();
auto grad_x_data = grad_x.DATA_PTR<scalar_t>();
for (ptrdiff_t e = 0; e < E; e++) {
for (ptrdiff_t m_out = 0; m_out < M_out; m_out++) {
......@@ -78,11 +80,11 @@ at::Tensor weighting_bw_w(at::Tensor grad_out, at::Tensor x, at::Tensor basis,
auto grad_weight = at::zeros({K, M_in, M_out}, grad_out.options());
AT_DISPATCH_FLOATING_TYPES(grad_out.scalar_type(), "weighting_bw_w", [&] {
auto grad_out_data = grad_out.data<scalar_t>();
auto x_data = x.data<scalar_t>();
auto basis_data = basis.data<scalar_t>();
auto weight_index_data = weight_index.data<int64_t>();
auto grad_weight_data = grad_weight.data<scalar_t>();
auto grad_out_data = grad_out.DATA_PTR<scalar_t>();
auto x_data = x.DATA_PTR<scalar_t>();
auto basis_data = basis.DATA_PTR<scalar_t>();
auto weight_index_data = weight_index.DATA_PTR<int64_t>();
auto grad_weight_data = grad_weight.DATA_PTR<scalar_t>();
for (ptrdiff_t e = 0; e < E; e++) {
for (ptrdiff_t m_out = 0; m_out < M_out; m_out++) {
......@@ -110,11 +112,11 @@ at::Tensor weighting_bw_b(at::Tensor grad_out, at::Tensor x, at::Tensor weight,
auto grad_basis = at::zeros({E, S}, grad_out.options());
AT_DISPATCH_FLOATING_TYPES(grad_out.scalar_type(), "weighting_bw_b", [&] {
auto grad_out_data = grad_out.data<scalar_t>();
auto x_data = x.data<scalar_t>();
auto weight_data = weight.data<scalar_t>();
auto weight_index_data = weight_index.data<int64_t>();
auto grad_basis_data = grad_basis.data<scalar_t>();
auto grad_out_data = grad_out.DATA_PTR<scalar_t>();
auto x_data = x.DATA_PTR<scalar_t>();
auto weight_data = weight.DATA_PTR<scalar_t>();
auto weight_index_data = weight_index.DATA_PTR<int64_t>();
auto grad_basis_data = grad_basis.DATA_PTR<scalar_t>();
for (ptrdiff_t e = 0; e < E; e++) {
for (ptrdiff_t m_out = 0; m_out < M_out; m_out++) {
......
cuda/basis_kernel.cu
View file @ 676aa072
......@@ -2,6 +2,8 @@
#include <ATen/cuda/detail/IndexUtils.cuh>
#include <ATen/cuda/detail/TensorInfo.cuh>
#include "compat.cuh"
#define THREADS 1024
#define BLOCKS(N) (N + THREADS - 1) / THREADS
......@@ -45,8 +47,8 @@ template <typename scalar_t> struct BasisForward {
at::cuda::detail::getTensorInfo<scalar_t, int64_t>(basis), \
at::cuda::detail::getTensorInfo<int64_t, int64_t>(weight_index), \
at::cuda::detail::getTensorInfo<scalar_t, int64_t>(PSEUDO), \
KERNEL_SIZE.data<int64_t>(), IS_OPEN_SPLINE.data<uint8_t>(), \
basis.numel()); \
KERNEL_SIZE.DATA_PTR<int64_t>(), \
IS_OPEN_SPLINE.DATA_PTR<uint8_t>(), basis.numel()); \
}); \
\
return std::make_tuple(basis, weight_index); \
......@@ -176,8 +178,8 @@ template <typename scalar_t> struct BasisBackward {
at::cuda::detail::getTensorInfo<scalar_t, int64_t>(grad_pseudo), \
at::cuda::detail::getTensorInfo<scalar_t, int64_t>(GRAD_BASIS), \
at::cuda::detail::getTensorInfo<scalar_t, int64_t>(PSEUDO), \
KERNEL_SIZE.data<int64_t>(), IS_OPEN_SPLINE.data<uint8_t>(), \
grad_pseudo.numel()); \
KERNEL_SIZE.DATA_PTR<int64_t>(), \
IS_OPEN_SPLINE.DATA_PTR<uint8_t>(), grad_pseudo.numel()); \
}); \
\
return grad_pseudo; \
......
cuda/compat.cuh 0 → 100644
View file @ 676aa072
#ifdef VERSION_GE_1_3
#define DATA_PTR data_ptr
#else
#define DATA_PTR data
#endif
setup.py
View file @ 676aa072
......@@ -2,21 +2,32 @@ from setuptools import setup, find_packages
import torch
from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])
extra_compile_args = []
if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
extra_compile_args += ['-DVERSION_GE_1_3']
ext_modules = [
CppExtension('torch_spline_conv.basis_cpu', ['cpu/basis.cpp']),
CppExtension('torch_spline_conv.weighting_cpu', ['cpu/weighting.cpp']),
CppExtension('torch_spline_conv.basis_cpu', ['cpu/basis.cpp'],
extra_compile_args=extra_compile_args),
CppExtension('torch_spline_conv.weighting_cpu', ['cpu/weighting.cpp'],
extra_compile_args=extra_compile_args),
]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
if CUDA_HOME is not None:
ext_modules += [
CUDAExtension('torch_spline_conv.basis_cuda',
['cuda/basis.cpp', 'cuda/basis_kernel.cu']),
['cuda/basis.cpp', 'cuda/basis_kernel.cu'],
extra_compile_args=extra_compile_args),
CUDAExtension('torch_spline_conv.weighting_cuda',
['cuda/weighting.cpp', 'cuda/weighting_kernel.cu']),
['cuda/weighting.cpp', 'cuda/weighting_kernel.cu'],
extra_compile_args=extra_compile_args),
]
__version__ = '1.1.0'
__version__ = '1.1.1'
url = 'https://github.com/rusty1s/pytorch_spline_conv'
install_requires = []
......
test/test_weighting.py
View file @ 676aa072
......@@ -32,7 +32,7 @@ def test_spline_weighting_forward(test, dtype, device):
@pytest.mark.parametrize('device', devices)
def test_spline_basis_backward(device):
def test_spline_weighting_backward(device):
pseudo = torch.rand((4, 2), dtype=torch.double, device=device)
kernel_size = tensor([5, 5], torch.long, device)
is_open_spline = tensor([1, 1], torch.uint8, device)
......
torch_spline_conv/__init__.py
View file @ 676aa072
......@@ -2,6 +2,6 @@ from .basis import SplineBasis
from .weighting import SplineWeighting
from .conv import SplineConv
__version__ = '1.1.0'
__version__ = '1.1.1'
__all__ = ['SplineBasis', 'SplineWeighting', 'SplineConv', '__version__']
torch_spline_conv/basis.py
View file @ 676aa072
......@@ -18,7 +18,8 @@ class SplineBasis(torch.autograd.Function):
@staticmethod
def forward(ctx, pseudo, kernel_size, is_open_spline, degree):
ctx.save_for_backward(pseudo)
ctx.kernel_size, ctx.is_open_spline = kernel_size, is_open_spline
ctx.kernel_size = kernel_size
ctx.is_open_spline = is_open_spline
ctx.degree = degree
op = get_func('{}_fw'.format(implemented_degrees[degree]), pseudo)
......
torch_spline_conv/conv.py
View file @ 676aa072
......@@ -38,18 +38,9 @@ class SplineConv(object):
:rtype: :class:`Tensor`
"""
@staticmethod
def apply(x,
edge_index,
pseudo,
weight,
kernel_size,
is_open_spline,
degree=1,
norm=True,
root_weight=None,
bias=None):
def apply(x, edge_index, pseudo, weight, kernel_size, is_open_spline,
degree=1, norm=True, root_weight=None, bias=None):
x = x.unsqueeze(-1) if x.dim() == 1 else x
pseudo = pseudo.unsqueeze(-1) if pseudo.dim() == 1 else pseudo
......@@ -58,8 +49,10 @@ class SplineConv(object):
n, m_out = x.size(0), weight.size(2)
# Weight each node.
data = SplineBasis.apply(pseudo, kernel_size, is_open_spline, degree)
out = SplineWeighting.apply(x[col], weight, *data)
basis, weight_index = SplineBasis.apply(pseudo, kernel_size,
is_open_spline, degree)
weight_index = weight_index.detach()
out = SplineWeighting.apply(x[col], weight, basis, weight_index)
# Convert e x m_out to n x m_out features.
row_expand = row.unsqueeze(-1).expand_as(out)
......
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