class SplineConv(MessagePassing):
r"""The spline-based convolutional operator from the `"SplineCNN: Fast
Geometric Deep Learning with Continuous B-Spline Kernels"
<https://arxiv.org/abs/1711.08920>`_ paper
.. math::
\mathbf{x}^{\prime}_i = \frac{1}{|\mathcal{N}(i)|} \sum_{j \in
\mathcal{N}(i)} \mathbf{x}_j \cdot
h_{\mathbf{\Theta}}(\mathbf{e}_{i,j}),
where :math:`h_{\mathbf{\Theta}}` denotes a kernel function defined
over the weighted B-Spline tensor product basis.
.. note::
Pseudo-coordinates must lay in the fixed interval :math:`[0, 1]` for
this method to work as intended.
Args:
in_channels (int or tuple): Size of each input sample, or :obj:`-1` to
derive the size from the first input(s) to the forward method.
A tuple corresponds to the sizes of source and target
dimensionalities.
out_channels (int): Size of each output sample.
dim (int): Pseudo-coordinate dimensionality.
kernel_size (int or [int]): Size of the convolving kernel.
is_open_spline (bool or [bool], optional): If set to :obj:`False`, the
operator will use a closed B-spline basis in this dimension.
(default :obj:`True`)
degree (int, optional): B-spline basis degrees. (default: :obj:`1`)
aggr (string, optional): The aggregation operator to use
(:obj:`"add"`, :obj:`"mean"`, :obj:`"max"`).
(default: :obj:`"mean"`)
root_weight (bool, optional): If set to :obj:`False`, the layer will
not add transformed root node features to the output.
(default: :obj:`True`)
bias (bool, optional): If set to :obj:`False`, the layer will not learn
an additive bias. (default: :obj:`True`)
**kwargs (optional): Additional arguments of
:class:`torch_geometric.nn.conv.MessagePassing`.
"""
def __init__(
self,
in_channels: Union[int, Tuple[int, int]],
out_channels: int,
dim: int,
kernel_size: Union[int, List[int]],
is_open_spline: bool = True,
degree: int = 1,
aggr: str = 'mean',
root_weight: bool = True,
bias: bool = True,
**kwargs,
):
super().__init__(aggr=aggr, **kwargs)
if spline_basis is None:
raise ImportError("'SplineConv' requires 'torch-spline-conv'")
self.in_channels = in_channels
self.out_channels = out_channels
self.dim = dim
self.degree = degree
self.root_weight = root_weight
kernel_size = torch.tensor(repeat(kernel_size, dim), dtype=torch.long)
self.register_buffer('kernel_size', kernel_size)
is_open_spline = repeat(is_open_spline, dim)
is_open_spline = torch.tensor(is_open_spline, dtype=torch.uint8)
self.register_buffer('is_open_spline', is_open_spline)
if isinstance(in_channels, int):
in_channels = (in_channels, in_channels)
self.K = kernel_size.prod().item()
if in_channels[0] > 0:
self.weight = Parameter(
torch.Tensor(self.K, in_channels[0], out_channels))
else:
self.weight = torch.nn.parameter.UninitializedParameter()
self._hook = self.register_forward_pre_hook(
self.initialize_parameters)
if root_weight:
self.lin = Linear(in_channels[1], out_channels, bias=False,
weight_initializer='uniform')
if bias:
self.bias = Parameter(torch.Tensor(out_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
if not isinstance(self.weight, nn.UninitializedParameter):
size = self.weight.size(0) * self.weight.size(1)
uniform(size, self.weight)
if self.root_weight:
self.lin.reset_parameters()
zeros(self.bias)
def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
edge_attr: OptTensor = None, size: Size = None) -> Tensor:
""""""
if isinstance(x, Tensor):
x: OptPairTensor = (x, x)
if not x[0].is_cuda:
warnings.warn(
'We do not recommend using the non-optimized CPU version of '
'`SplineConv`. If possible, please move your data to GPU.')
# propagate_type: (x: OptPairTensor, edge_attr: OptTensor)
out = self.propagate(edge_index, x=x, edge_attr=edge_attr, size=size)
x_r = x[1]
if x_r is not None and self.root_weight:
out += self.lin(x_r)
if self.bias is not None:
out += self.bias
return out
def message(self, x_j: Tensor, edge_attr: Tensor) -> Tensor:
data = spline_basis(edge_attr, self.kernel_size, self.is_open_spline,
self.degree)
return spline_weighting(x_j, self.weight, *data)
@torch.no_grad()
def initialize_parameters(self, module, input):
if isinstance(self.weight, torch.nn.parameter.UninitializedParameter):
x = input[0][0] if isinstance(input, tuple) else input[0]
in_channels = x.size(-1)
self.weight.materialize((self.K, in_channels, self.out_channels))
size = self.weight.size(0) * self.weight.size(1)
uniform(size, self.weight)
module._hook.remove()
delattr(module, '_hook')
def __repr__(self) -> str:
return (f'{self.__class__.__name__}({self.in_channels}, '
f'{self.out_channels}, dim={self.dim})')
class GMMConv(MessagePassing):
r"""The gaussian mixture model convolutional operator from the `"Geometric
Deep Learning on Graphs and Manifolds using Mixture Model CNNs"
<https://arxiv.org/abs/1611.08402>`_ paper
.. math::
\mathbf{x}^{\prime}_i = \frac{1}{|\mathcal{N}(i)|}
\sum_{j \in \mathcal{N}(i)} \frac{1}{K} \sum_{k=1}^K
\mathbf{w}_k(\mathbf{e}_{i,j}) \odot \mathbf{\Theta}_k \mathbf{x}_j,
where
.. math::
\mathbf{w}_k(\mathbf{e}) = \exp \left( -\frac{1}{2} {\left(
\mathbf{e} - \mathbf{\mu}_k \right)}^{\top} \Sigma_k^{-1}
\left( \mathbf{e} - \mathbf{\mu}_k \right) \right)
denotes a weighting function based on trainable mean vector
:math:`\mathbf{\mu}_k` and diagonal covariance matrix
:math:`\mathbf{\Sigma}_k`.
.. note::
The edge attribute :math:`\mathbf{e}_{ij}` is usually given by
:math:`\mathbf{e}_{ij} = \mathbf{p}_j - \mathbf{p}_i`, where
:math:`\mathbf{p}_i` denotes the position of node :math:`i` (see
:class:`torch_geometric.transform.Cartesian`).
Args:
in_channels (int or tuple): Size of each input sample, or :obj:`-1` to
derive the size from the first input(s) to the forward method.
A tuple corresponds to the sizes of source and target
dimensionalities.
out_channels (int): Size of each output sample.
dim (int): Pseudo-coordinate dimensionality.
kernel_size (int): Number of kernels :math:`K`.
separate_gaussians (bool, optional): If set to :obj:`True`, will
learn separate GMMs for every pair of input and output channel,
inspired by traditional CNNs. (default: :obj:`False`)
aggr (string, optional): The aggregation operator to use
(:obj:`"add"`, :obj:`"mean"`, :obj:`"max"`).
(default: :obj:`"mean"`)
root_weight (bool, optional): If set to :obj:`False`, the layer will
not add transformed root node features to the output.
(default: :obj:`True`)
bias (bool, optional): If set to :obj:`False`, the layer will not learn
an additive bias. (default: :obj:`True`)
**kwargs (optional): Additional arguments of
:class:`torch_geometric.nn.conv.MessagePassing`.
"""
def __init__(self,
in_channels: Union[int, Tuple[int, int]],
out_channels: int,
dim: int,
kernel_size: int,
separate_gaussians: bool = False,
aggr: str = 'mean',
root_weight: bool = True,
bias: bool = True,
**kwargs):
super(GMMConv, self).__init__(aggr=aggr, **kwargs)
self.in_channels = in_channels
self.out_channels = out_channels
self.dim = dim
self.kernel_size = kernel_size
self.separate_gaussians = separate_gaussians
self.root_weight = root_weight
if isinstance(in_channels, int):
in_channels = (in_channels, in_channels)
self.rel_in_channels = in_channels[0]
if in_channels[0] > 0:
self.g = Parameter(
Tensor(in_channels[0], out_channels * kernel_size))
if not self.separate_gaussians:
self.mu = Parameter(Tensor(kernel_size, dim))
self.sigma = Parameter(Tensor(kernel_size, dim))
if self.separate_gaussians:
self.mu = Parameter(
Tensor(in_channels[0], out_channels, kernel_size, dim))
self.sigma = Parameter(
Tensor(in_channels[0], out_channels, kernel_size, dim))
else:
self.g = torch.nn.parameter.UninitializedParameter()
self.mu = torch.nn.parameter.UninitializedParameter()
self.sigma = torch.nn.parameter.UninitializedParameter()
self._hook = self.register_forward_pre_hook(
self.initialize_parameters)
if root_weight:
self.root = Linear(in_channels[1],
out_channels,
bias=False,
weight_initializer='glorot')
if bias:
self.bias = Parameter(torch.Tensor(out_channels))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
if not isinstance(self.g, torch.nn.UninitializedParameter):
glorot(self.g)
glorot(self.mu)
glorot(self.sigma)
if self.root_weight:
self.root.reset_parameters()
zeros(self.bias)
def forward(self,
x: Union[Tensor, OptPairTensor],
edge_index: Adj,
edge_attr: OptTensor = None,
size: Size = None):
""""""
if isinstance(x, Tensor):
x: OptPairTensor = (x, x)
# propagate_type: (x: OptPairTensor, edge_attr: OptTensor)
if not self.separate_gaussians:
out: OptPairTensor = (torch.matmul(x[0], self.g), x[1])
out = self.propagate(edge_index,
x=out,
edge_attr=edge_attr,
size=size)
else:
out = self.propagate(edge_index,
x=x,
edge_attr=edge_attr,
size=size)
x_r = x[1]
if x_r is not None and self.root is not None:
out += self.root(x_r)
if self.bias is not None:
out += self.bias
return out
def message(self, x_j: Tensor, edge_attr: Tensor):
EPS = 1e-15
F, M = self.rel_in_channels, self.out_channels
(E, D), K = edge_attr.size(), self.kernel_size
if not self.separate_gaussians:
gaussian = -0.5 * (edge_attr.view(E, 1, D) -
self.mu.view(1, K, D)).pow(2)
gaussian = gaussian / (EPS + self.sigma.view(1, K, D).pow(2))
gaussian = torch.exp(gaussian.sum(dim=-1)) # [E, K]
return (x_j.view(E, K, M) * gaussian.view(E, K, 1)).sum(dim=-2)
else:
gaussian = -0.5 * (edge_attr.view(E, 1, 1, 1, D) -
self.mu.view(1, F, M, K, D)).pow(2)
gaussian = gaussian / (EPS + self.sigma.view(1, F, M, K, D).pow(2))
gaussian = torch.exp(gaussian.sum(dim=-1)) # [E, F, M, K]
gaussian = gaussian * self.g.view(1, F, M, K)
gaussian = gaussian.sum(dim=-1) # [E, F, M]
return (x_j.view(E, F, 1) * gaussian).sum(dim=-2) # [E, M]
@torch.no_grad()
def initialize_parameters(self, module, input):
if isinstance(self.g, torch.nn.parameter.UninitializedParameter):
x = input[0][0] if isinstance(input, tuple) else input[0]
in_channels = x.size(-1)
out_channels, kernel_size = self.out_channels, self.kernel_size
self.g.materialize((in_channels, out_channels * kernel_size))
if not self.separate_gaussians:
self.mu.materialize((kernel_size, self.dim))
self.sigma.materialize((kernel_size, self.dim))
else:
self.mu.materialize(
(in_channels, out_channels, kernel_size, self.dim))
self.sigma.materialize(
(in_channels, out_channels, kernel_size, self.dim))
glorot(self.g)
glorot(self.mu)
glorot(self.sigma)
module._hook.remove()
delattr(module, '_hook')
def __repr__(self) -> str:
return '{}({}, {}, dim={})'.format(self.__class__.__name__,
self.in_channels, self.out_channels,
self.dim)
class ARMAConv(MessagePassing):
r"""The ARMA graph convolutional operator from the `"Graph Neural Networks
with Convolutional ARMA Filters" <https://arxiv.org/abs/1901.01343>`_ paper
.. math::
\mathbf{X}^{\prime} = \frac{1}{K} \sum_{k=1}^K \mathbf{X}_k^{(T)},
with :math:`\mathbf{X}_k^{(T)}` being recursively defined by
.. math::
\mathbf{X}_k^{(t+1)} = \sigma \left( \mathbf{\hat{L}}
\mathbf{X}_k^{(t)} \mathbf{W} + \mathbf{X}^{(0)} \mathbf{V} \right),
where :math:`\mathbf{\hat{L}} = \mathbf{I} - \mathbf{L} = \mathbf{D}^{-1/2}
\mathbf{A} \mathbf{D}^{-1/2}` denotes the
modified Laplacian :math:`\mathbf{L} = \mathbf{I} - \mathbf{D}^{-1/2}
\mathbf{A} \mathbf{D}^{-1/2}`.
Args:
in_channels (int): Size of each input sample, or :obj:`-1` to derive
the size from the first input(s) to the forward method.
out_channels (int): Size of each output sample
:math:`\mathbf{x}^{(t+1)}`.
num_stacks (int, optional): Number of parallel stacks :math:`K`.
(default: :obj:`1`).
num_layers (int, optional): Number of layers :math:`T`.
(default: :obj:`1`)
act (callable, optional): Activation function :math:`\sigma`.
(default: :meth:`torch.nn.ReLU()`)
shared_weights (int, optional): If set to :obj:`True` the layers in
each stack will share the same parameters. (default: :obj:`False`)
dropout (float, optional): Dropout probability of the skip connection.
(default: :obj:`0.`)
bias (bool, optional): If set to :obj:`False`, the layer will not learn
an additive bias. (default: :obj:`True`)
**kwargs (optional): Additional arguments of
:class:`torch_geometric.nn.conv.MessagePassing`.
"""
def __init__(self,
in_channels: int,
out_channels: int,
num_stacks: int = 1,
num_layers: int = 1,
shared_weights: bool = False,
act: Optional[Callable] = ReLU(),
dropout: float = 0.,
bias: bool = True,
**kwargs):
kwargs.setdefault('aggr', 'add')
super(ARMAConv, self).__init__(**kwargs)
self.in_channels = in_channels
self.out_channels = out_channels
self.num_stacks = num_stacks
self.num_layers = num_layers
self.act = act
self.shared_weights = shared_weights
self.dropout = dropout
K, T, F_in, F_out = num_stacks, num_layers, in_channels, out_channels
T = 1 if self.shared_weights else T
self.weight = Parameter(torch.Tensor(max(1, T - 1), K, F_out, F_out))
if in_channels > 0:
self.init_weight = Parameter(torch.Tensor(K, F_in, F_out))
self.root_weight = Parameter(torch.Tensor(T, K, F_in, F_out))
else:
self.init_weight = torch.nn.parameter.UninitializedParameter()
self.root_weight = torch.nn.parameter.UninitializedParameter()
self._hook = self.register_forward_pre_hook(
self.initialize_parameters)
if bias:
self.bias = Parameter(torch.Tensor(T, K, 1, F_out))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self):
glorot(self.weight)
if not isinstance(self.init_weight, torch.nn.UninitializedParameter):
glorot(self.init_weight)
glorot(self.root_weight)
zeros(self.bias)
def forward(self,
x: Tensor,
edge_index: Adj,
edge_weight: OptTensor = None) -> Tensor:
""""""
if isinstance(edge_index, Tensor):
edge_index, edge_weight = gcn_norm( # yapf: disable
edge_index,
edge_weight,
x.size(self.node_dim),
add_self_loops=False,
dtype=x.dtype)
elif isinstance(edge_index, SparseTensor):
edge_index = gcn_norm( # yapf: disable
edge_index,
edge_weight,
x.size(self.node_dim),
add_self_loops=False,
dtype=x.dtype)
x = x.unsqueeze(-3)
out = x
for t in range(self.num_layers):
if t == 0:
out = out @ self.init_weight
else:
out = out @ self.weight[0 if self.shared_weights else t - 1]
# propagate_type: (x: Tensor, edge_weight: OptTensor)
out = self.propagate(edge_index,
x=out,
edge_weight=edge_weight,
size=None)
root = F.dropout(x, p=self.dropout, training=self.training)
out += root @ self.root_weight[0 if self.shared_weights else t]
if self.bias is not None:
out += self.bias[0 if self.shared_weights else t]
if self.act is not None:
out = self.act(out)
return out.mean(dim=-3)
def message(self, x_j: Tensor, edge_weight: Tensor) -> Tensor:
return edge_weight.view(-1, 1) * x_j
def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
return matmul(adj_t, x, reduce=self.aggr)
@torch.no_grad()
def initialize_parameters(self, module, input):
if isinstance(self.init_weight, nn.parameter.UninitializedParameter):
F_in, F_out = input[0].size(-1), self.out_channels
T, K = self.weight.size(0) + 1, self.weight.size(1)
self.init_weight.materialize((K, F_in, F_out))
self.root_weight.materialize((T, K, F_in, F_out))
glorot(self.init_weight)
glorot(self.root_weight)
module._hook.remove()
delattr(module, '_hook')
def __repr__(self) -> str:
return '{}({}, {}, num_stacks={}, num_layers={})'.format(
self.__class__.__name__, self.in_channels, self.out_channels,
self.num_stacks, self.num_layers)
