utils.py 2.31 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
5
import torch
from torch.autograd import Function

from .._ext import ffi

rusty1s's avatar
rusty1s committed
6
implemented_degrees = {1: 'linear', 2: 'quadratic', 3: 'cubic'}
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
9
10
11
12
13
14
15

def get_func(name, tensor):
    typename = type(tensor).__name__.replace('Tensor', '')
    cuda = 'cuda_' if tensor.is_cuda else ''
    func = getattr(ffi, 'spline_{}_{}{}'.format(name, cuda, typename))
    return func


rusty1s's avatar
rusty1s committed
16
def spline_basis(degree, pseudo, kernel_size, is_open_spline, K):
rusty1s's avatar
rusty1s committed
17
18
19
20
21
    s = (degree + 1)**kernel_size.size(0)
    pseudo = pseudo.unsqueeze(-1) if pseudo.dim() == 1 else pseudo
    basis = pseudo.new(pseudo.size(0), s)
    weight_index = kernel_size.new(pseudo.size(0), s)

rusty1s's avatar
rusty1s committed
22
    degree = implemented_degrees.get(degree)
rusty1s's avatar
rusty1s committed
23
24
25
26
    if degree is None:
        raise NotImplementedError('Basis computation not implemented for '
                                  'specified B-spline degree')

rusty1s's avatar
rusty1s committed
27
    func = get_func('basis_{}'.format(degree), pseudo)
rusty1s's avatar
rusty1s committed
28
29
    func(basis, weight_index, pseudo, kernel_size, is_open_spline, K)
    return basis, weight_index
rusty1s's avatar
rusty1s committed
30
31


rusty1s's avatar
rusty1s committed
32
33
34
35
36
def spline_weighting_fw(x, weight, basis, weight_index):
    output = x.new(x.size(0), weight.size(2))
    func = get_func('spline_weighting_fw', x)
    func(output, x, weight, basis, weight_index)
    return output
rusty1s's avatar
rusty1s committed
37
38


rusty1s's avatar
rusty1s committed
39
40
41
42
43
44
def spline_weighting_bw(grad_output, x, weight, basis, weight_index):
    grad_input = x.new(x.size(0), weight.size(1))
    grad_weight = x.new(weight)
    func = get_func('spline_weighting_bw', x)
    func(grad_input, grad_weight, grad_output, x, weight, basis, weight_index)
    return grad_input, grad_weight
rusty1s's avatar
rusty1s committed
45
46
47
48
49
50
51
52
53


class SplineWeighting(Function):
    def __init__(self, basis, weight_index):
        super(SplineWeighting, self).__init__()
        self.basis = basis
        self.weight_index = weight_index

    def forward(self, x, weight):
rusty1s's avatar
rusty1s committed
54
55
56
        self.save_for_backward(x, weight)
        basis, weight_index = self.basis, self.weight_index
        return spline_weighting_fw(x, weight, basis, weight_index)
rusty1s's avatar
rusty1s committed
57
58

    def backward(self, grad_output):
rusty1s's avatar
rusty1s committed
59
60
61
        x, weight = self.saved_tensors
        basis, weight_index = self.basis, self.weight_index
        return spline_weighting_bw(grad_output, x, weight, basis, weight_index)
rusty1s's avatar
rusty1s committed
62
63
64
65


def spline_weighting(x, weight, basis, weight_index):
    if torch.is_tensor(x):
rusty1s's avatar
rusty1s committed
66
        return spline_weighting_fw(x, weight, basis, weight_index)
rusty1s's avatar
rusty1s committed
67
68
    else:
        return SplineWeighting(basis, weight_index)(x, weight)