conv.py 3.05 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
rusty1s's avatar
linting  
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
4
from .basis import SplineBasis
from .weighting import SplineWeighting
rusty1s's avatar
linting  
rusty1s committed
5

rusty1s's avatar
rusty1s committed
6
from .utils.degree import degree as node_degree
rusty1s's avatar
rusty1s committed
7
8


rusty1s's avatar
rusty1s committed
9
class SplineConv(object):
rusty1s's avatar
typo  
rusty1s committed
10
    """Applies the spline-based convolution operator :math:`(f \star g)(i) =
rusty1s's avatar
rusty1s committed
11
12
    \frac{1}{|\mathcal{N}(i)|} \sum_{l=1}^{M_{in}} \sum_{j \in \mathcal{N}(i)}
    f_l(j) \cdot g_l(u(i, j))` over several node features of an input graph.
rusty1s's avatar
typo  
rusty1s committed
13
14
    The kernel function :math:`g_l` is defined over the weighted B-spline
    tensor product basis for a single input feature map :math:`l`.
rusty1s's avatar
rusty1s committed
15
16

    Args:
rusty1s's avatar
rusty1s committed
17
        x (:class:`Tensor`): Input node features of shape
rusty1s's avatar
rusty1s committed
18
19
20
21
22
23
24
25
26
27
28
29
            (number_of_nodes x in_channels).
        edge_index (:class:`LongTensor`): Graph edges, given by source and
            target indices, of shape (2 x number_of_edges) in the fixed
            interval [0, 1].
        pseudo (:class:`Tensor`): Edge attributes, ie. pseudo coordinates,
            of shape (number_of_edges x number_of_edge_attributes).
        weight (:class:`Tensor`): Trainable weight parameters of shape
            (kernel_size x in_channels x out_channels).
        kernel_size (:class:`LongTensor`): Number of trainable weight
            parameters in each edge dimension.
        is_open_spline (:class:`ByteTensor`): Whether to use open or closed
            B-spline bases for each dimension.
rusty1s's avatar
rusty1s committed
30
        degree (int, optional): B-spline basis degree. (default: :obj:`1`)
rusty1s's avatar
rusty1s committed
31
32
        norm (bool, optional): Whether to normalize output by node degree.
            (default: :obj:`True`)
rusty1s's avatar
rusty1s committed
33
        root_weight (:class:`Tensor`, optional): Additional shared trainable
rusty1s's avatar
rusty1s committed
34
            parameters for each feature of the root node of shape
rusty1s's avatar
rusty1s committed
35
36
37
38
39
            (in_channels x out_channels). (default: :obj:`None`)
        bias (:class:`Tensor`, optional): Optional bias of shape
            (out_channels). (default: :obj:`None`)

    :rtype: :class:`Tensor`
rusty1s's avatar
rusty1s committed
40
    """
rusty1s's avatar
rusty1s committed
41

rusty1s's avatar
rusty1s committed
42
    @staticmethod
rusty1s's avatar
rusty1s committed
43
    def apply(x,
rusty1s's avatar
rusty1s committed
44
45
46
47
48
49
              edge_index,
              pseudo,
              weight,
              kernel_size,
              is_open_spline,
              degree=1,
rusty1s's avatar
rusty1s committed
50
              norm=True,
rusty1s's avatar
rusty1s committed
51
52
              root_weight=None,
              bias=None):
rusty1s's avatar
rusty1s committed
53

rusty1s's avatar
rusty1s committed
54
        x = x.unsqueeze(-1) if x.dim() == 1 else x
rusty1s's avatar
rusty1s committed
55
        pseudo = pseudo.unsqueeze(-1) if pseudo.dim() == 1 else pseudo
rusty1s's avatar
rusty1s committed
56

rusty1s's avatar
rusty1s committed
57
        row, col = edge_index
rusty1s's avatar
rusty1s committed
58
        n, m_out = x.size(0), weight.size(2)
rusty1s's avatar
rusty1s committed
59

rusty1s's avatar
rusty1s committed
60
        # Weight each node.
rusty1s's avatar
rusty1s committed
61
        data = SplineBasis.apply(pseudo, kernel_size, is_open_spline, degree)
rusty1s's avatar
rusty1s committed
62
        out = SplineWeighting.apply(x[col], weight, *data)
rusty1s's avatar
rusty1s committed
63

rusty1s's avatar
rusty1s committed
64
        # Convert e x m_out to n x m_out features.
rusty1s's avatar
rename  
rusty1s committed
65
        row_expand = row.unsqueeze(-1).expand_as(out)
rusty1s's avatar
rusty1s committed
66
        out = x.new_zeros((n, m_out)).scatter_add_(0, row_expand, out)
rusty1s's avatar
rusty1s committed
67

rusty1s's avatar
rusty1s committed
68
69
70
71
        # Normalize out by node degree (if wished).
        if norm:
            deg = node_degree(row, n, out.dtype, out.device)
            out = out / deg.unsqueeze(-1).clamp(min=1)
rusty1s's avatar
rusty1s committed
72

rusty1s's avatar
rusty1s committed
73
74
        # Weight root node separately (if wished).
        if root_weight is not None:
rusty1s's avatar
rusty1s committed
75
            out = out + torch.mm(x, root_weight)
rusty1s's avatar
rusty1s committed
76

rusty1s's avatar
rusty1s committed
77
78
        # Add bias (if wished).
        if bias is not None:
rusty1s's avatar
rusty1s committed
79
            out = out + bias
rusty1s's avatar
rusty1s committed
80

rusty1s's avatar
rename  
rusty1s committed
81
        return out