weighting.py

import torch
import weighting_cpu

if torch.cuda.is_available():
    import weighting_cuda


def get_func(name, tensor):
    module = weighting_cuda if tensor.is_cuda else weighting_cpu
    return getattr(module, name)


class SplineWeighting(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, weight, basis, weight_index):
        ctx.weight_index = weight_index
        ctx.save_for_backward(x, weight, basis)
        op = get_func('weighting_fw', x)
        out = op(x, weight, basis, weight_index)
        return out

    @staticmethod
    def backward(ctx, grad_out):
        x, weight, basis = ctx.saved_tensors
        grad_x = grad_weight = grad_basis = None

        if ctx.needs_input_grad[0]:
            op = get_func('weighting_bw_x', x)
            grad_x = op(grad_out, weight, basis, ctx.weight_index)

        if ctx.needs_input_grad[1]:
            op = get_func('weighting_bw_w', x)
            grad_weight = op(grad_out, x, basis, ctx.weight_index,
                             weight.size(0))

        if ctx.needs_input_grad[2]:
            op = get_func('weighting_bw_b', x)
            grad_basis = op(grad_out, x, weight, ctx.weight_index)

        return grad_x, grad_weight, grad_basis, None