ffi.py 2.52 KB
Newer Older
1
2
3
4
5
from .._ext import ffi as ext

implemented_degrees = {1: 'linear', 2: 'quadratic', 3: 'cubic'}


rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
12
def get_degree_str(degree):
    degree = implemented_degrees.get(degree)
    assert degree is not None, (
        'No implementation found for specified B-spline degree')
    return degree


13
14
15
16
17
18
19
20
21
22
23
def get_func(name, tensor):
    typename = type(tensor).__name__.replace('Tensor', '')
    cuda = 'cuda_' if tensor.is_cuda else ''
    func = getattr(ext, 'spline_{}_{}{}'.format(name, cuda, typename))
    return func


def spline_basis_forward(degree, pseudo, kernel_size, is_open_spline, K):
    s = (degree + 1)**kernel_size.size(0)
    basis = pseudo.new(pseudo.size(0), s)
    weight_index = kernel_size.new(pseudo.size(0), s)
rusty1s's avatar
rusty1s committed
24
    func = get_func('{}_basis_forward'.format(get_degree_str(degree)), pseudo)
25
26
27
28
    func(basis, weight_index, pseudo, kernel_size, is_open_spline, K)
    return basis, weight_index


rusty1s's avatar
rusty1s committed
29
def spline_basis_backward(degree, grad_basis, pseudo, kernel_size,
rusty1s's avatar
rusty1s committed
30
                          is_open_spline):  # pragma: no cover
rusty1s's avatar
rusty1s committed
31
32
33
34
35
36
    grad_pseudo = pseudo.new(pseudo.size())
    func = get_func('{}_basis_backward'.format(get_degree_str(degree)), pseudo)
    func(grad_pseudo, grad_basis, pseudo, kernel_size, is_open_spline)
    return grad_pseudo


37
38
39
40
41
42
43
def spline_weighting_forward(x, weight, basis, weight_index):
    output = x.new(x.size(0), weight.size(2))
    func = get_func('weighting_forward', x)
    func(output, x, weight, basis, weight_index)
    return output


rusty1s's avatar
rusty1s committed
44
45
def spline_weighting_backward_input(grad_output, weight, basis,
                                    weight_index):  # pragma: no cover
46
47
    grad_input = grad_output.new(grad_output.size(0), weight.size(1))
    func = get_func('weighting_backward_input', grad_output)
rusty1s's avatar
rusty1s committed
48
    weight = weight.transpose(1, 2).contiguous()
49
    func(grad_input, grad_output, weight, basis, weight_index)
rusty1s's avatar
rusty1s committed
50
    weight = weight.transpose(1, 2).contiguous()
51
52
53
    return grad_input


rusty1s's avatar
rusty1s committed
54
55
def spline_weighting_backward_basis(grad_output, x, weight,
                                    weight_index):  # pragma: no cover
rusty1s's avatar
rusty1s committed
56
    grad_basis = x.new(weight_index.size()).fill_(0)
57
58
59
    func = get_func('weighting_backward_basis', x)
    func(grad_basis, grad_output, x, weight, weight_index)
    return grad_basis
rusty1s's avatar
rusty1s committed
60
61


rusty1s's avatar
rusty1s committed
62
63
def spline_weighting_backward_weight(grad_output, x, basis, weight_index,
                                     K):  # pragma: no cover
rusty1s's avatar
rusty1s committed
64
65
66
67
    grad_weight = x.new(K, x.size(1), grad_output.size(1)).fill_(0)
    func = get_func('weighting_backward_weight', x)
    func(grad_weight, grad_output, x, basis, weight_index)
    return grad_weight