class LSTMAggregation(Aggregation):
r"""Performs LSTM-style aggregation in which the elements to aggregate are
interpreted as a sequence, as described in the `"Inductive Representation
Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper.
.. warning::
:class:`LSTMAggregation` is not a permutation-invariant operator.
Args:
in_channels (int): Size of each input sample.
out_channels (int): Size of each output sample.
**kwargs (optional): Additional arguments of :class:`torch.nn.LSTM`.
"""
def __init__(self, in_channels: int, out_channels: int, **kwargs):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.lstm = LSTM(in_channels, out_channels, batch_first=True, **kwargs)
self.reset_parameters()
def reset_parameters(self):
self.lstm.reset_parameters()
def forward(self,
x: Tensor,
index: Optional[Tensor] = None,
ptr: Optional[Tensor] = None,
dim_size: Optional[int] = None,
dim: int = -2) -> Tensor:
x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim)
return self.lstm(x)[0][:, -1]
def __repr__(self) -> str:
return (f'{self.__class__.__name__}({self.in_channels}, '
f'{self.out_channels})')
class DenseJK(nn.Module):
def __init__(self, mode, channels=None, num_layers=None):
super(DenseJK, self).__init__()
self.channel = channels
self.mode = mode.lower()
assert self.mode in ['cat', 'max', 'lstm']
if mode == 'lstm':
assert channels is not None
assert num_layers is not None
self.lstm = LSTM(channels,
channels * num_layers // 2,
bidirectional=True,
batch_first=True)
self.att = Linear(2 * channels * num_layers // 2, 1)
self.reset_parameters()
def reset_parameters(self):
if hasattr(self, 'lstm'):
self.lstm.reset_parameters()
if hasattr(self, 'att'):
self.att.reset_parameters()
def forward(self, xs):
r"""Aggregates representations across different layers.
Args:
xs [batch, nodes, featdim*3]
"""
xs = torch.split(xs, self.channel, -1) # list of batch, node, featdim
xs = torch.stack(xs, 2) #[batch, nodes, num_layers, num_channels]
shape = xs.shape
x = xs.reshape(
(-1, shape[2],
shape[3])) # [ngraph * num_nodes , num_layers, num_channels]
alpha, _ = self.lstm(x)
alpha = self.att(alpha).squeeze(-1) # [ngraph * num_nodes, num_layers]
alpha = torch.softmax(alpha, dim=-1)
x = (x * alpha.unsqueeze(-1)).sum(dim=1)
x = x.reshape((shape[0], shape[1], shape[3]))
return x
def __repr__(self):
return '{}({})'.format(self.__class__.__name__, self.mode)
class Stock_LSTM(nn.Module):
"""
I prefer using this Stock LSTM for numerical stability.
"""
def __init__(self, x, R, W, h, L, v_t):
super(Stock_LSTM, self).__init__()
self.x = x
self.R = R
self.W = W
self.h = h
self.L = L
self.v_t = v_t
self.LSTM = LSTM(input_size=self.x + self.R * self.W,
hidden_size=h,
num_layers=L,
batch_first=True,
dropout=0.1,
bidirectional=True)
self.last = nn.Linear(self.h * 2, self.v_t)
self.st = None
def forward(self, input_x):
"""
:param input_x: input and memory values
:return:
"""
assert (self.st is not None)
o, st = self.LSTM(input_x, self.st)
if (st[0] != st[0]).any():
with open("debug/lstm.pkl") as f:
pickle.dump(self, f)
with open("debug/lstm.pkl") as f:
pickle.dump(input_x, f)
raise ("LSTM produced a NAN, objects dumped.")
return self.last(o), st
def reset_parameters(self):
self.LSTM.reset_parameters()
self.last.reset_parameters()
def assign_states_tuple(self, states_tuple):
self.st = states_tuple
class LSTMAggregation(Aggregation):
r"""Performs LSTM-style aggregation in which the elements to aggregate are
interpreted as a sequence.
.. warn::
:class:`LSTMAggregation` is not permutation-invariant.
.. note::
:class:`LSTMAggregation` requires sorted indices.
Args:
in_channels (int): Size of each input sample.
out_channels (int): Size of each output sample.
**kwargs (optional): Additional arguments of :class:`torch.nn.LSTM`.
"""
def __init__(self, in_channels: int, out_channels: int, **kwargs):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.lstm = LSTM(in_channels, out_channels, batch_first=True, **kwargs)
self.reset_parameters()
def reset_parameters(self):
self.lstm.reset_parameters()
def forward(self,
x: Tensor,
index: Optional[Tensor] = None,
*,
ptr: Optional[Tensor] = None,
dim_size: Optional[int] = None,
dim: int = -2) -> Tensor:
x, _ = self.to_dense_batch(x, index, ptr, dim_size, dim)
return self.lstm(x)[0][:, -1]
def __repr__(self) -> str:
return (f'{self.__class__.__name__}({self.in_channels}, '
f'{self.out_channels})')
class JumpingKnowledge(torch.nn.Module):
r"""The Jumping Knowledge layer aggregation module from the
`"Representation Learning on Graphs with Jumping Knowledge Networks"
<https://arxiv.org/abs/1806.03536>`_ paper based on either
**concatenation** (:obj:`"cat"`)
.. math::
\mathbf{x}_v^{(1)} \, \Vert \, \ldots \, \Vert \, \mathbf{x}_v^{(T)}
**max pooling** (:obj:`"max"`)
.. math::
\max \left( \mathbf{x}_v^{(1)}, \ldots, \mathbf{x}_v^{(T)} \right)
or **weighted summation**
.. math::
\sum_{t=1}^T \alpha_v^{(t)} \mathbf{x}_v^{(t)}
with attention scores :math:`\alpha_v^{(t)}` obtained from a bi-directional
LSTM (:obj:`"lstm"`).
Args:
mode (string): The aggregation scheme to use
(:obj:`"cat"`, :obj:`"max"` or :obj:`"lstm"`).
channels (int, optional): The number of channels per representation.
Needs to be only set for LSTM-style aggregation.
(default: :obj:`None`)
num_layers (int, optional): The number of layers to aggregate. Needs to
be only set for LSTM-style aggregation. (default: :obj:`None`)
"""
def __init__(self, mode, channels=None, num_layers=None):
super(JumpingKnowledge, self).__init__()
self.mode = mode.lower()
assert self.mode in ['cat', 'max', 'lstm']
if mode == 'lstm':
assert channels is not None
assert num_layers is not None
self.lstm = LSTM(
channels,
channels * num_layers // 2,
bidirectional=True,
batch_first=True)
self.att = Linear(2 * channels * num_layers // 2, 1)
self.reset_parameters()
def reset_parameters(self):
if hasattr(self, 'lstm'):
self.lstm.reset_parameters()
if hasattr(self, 'att'):
self.att.reset_parameters()
def forward(self, xs):
r"""Aggregates representations across different layers.
Args:
xs (list or tuple): List containing layer-wise representations.
"""
assert isinstance(xs, list) or isinstance(xs, tuple)
if self.mode == 'cat':
return torch.cat(xs, dim=-1)
elif self.mode == 'max':
return torch.stack(xs, dim=-1).max(dim=-1)[0]
elif self.mode == 'lstm':
x = torch.stack(xs, dim=1) # [num_nodes, num_layers, num_channels]
alpha, _ = self.lstm(x)
alpha = self.att(alpha).squeeze(-1) # [num_nodes, num_layers]
alpha = torch.softmax(alpha, dim=-1)
return (x * alpha.unsqueeze(-1)).sum(dim=1)
def __repr__(self):
return '{}({})'.format(self.__class__.__name__, self.mode)
class SAGEConv(MessagePassing):
r"""The GraphSAGE operator from the `"Inductive Representation Learning on
Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper
.. math::
\mathbf{x}^{\prime}_i = \mathbf{W}_1 \mathbf{x}_i + \mathbf{W}_2 \cdot
\mathrm{mean}_{j \in \mathcal{N(i)}} \mathbf{x}_j
If :obj:`project = True`, then :math:`\mathbf{x}_j` will first get
projected via
.. math::
\mathbf{x}_j \leftarrow \sigma ( \mathbf{W}_3 \mathbf{x}_j +
\mathbf{b})
as described in Eq. (3) of the paper.
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.
aggr (string, optional): The aggregation scheme to use
(:obj:`"mean"`, :obj:`"max"`, :obj:`"lstm"`).
(default: :obj:`"add"`)
normalize (bool, optional): If set to :obj:`True`, output features
will be :math:`\ell_2`-normalized, *i.e.*,
:math:`\frac{\mathbf{x}^{\prime}_i}
{\| \mathbf{x}^{\prime}_i \|_2}`.
(default: :obj:`False`)
root_weight (bool, optional): If set to :obj:`False`, the layer will
not add transformed root node features to the output.
(default: :obj:`True`)
project (bool, optional): If set to :obj:`True`, the layer will apply a
linear transformation followed by an activation function before
aggregation (as described in Eq. (3) of the paper).
(default: :obj:`False`)
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`.
Shapes:
- **inputs:**
node features :math:`(|\mathcal{V}|, F_{in})` or
:math:`((|\mathcal{V_s}|, F_{s}), (|\mathcal{V_t}|, F_{t}))`
if bipartite,
edge indices :math:`(2, |\mathcal{E}|)`
- **outputs:** node features :math:`(|\mathcal{V}|, F_{out})` or
:math:`(|\mathcal{V_t}|, F_{out})` if bipartite
"""
def __init__(
self,
in_channels: Union[int, Tuple[int, int]],
out_channels: int,
aggr: str = 'mean',
normalize: bool = False,
root_weight: bool = True,
project: bool = False,
bias: bool = True,
**kwargs,
):
kwargs['aggr'] = aggr if aggr != 'lstm' else None
super().__init__(**kwargs)
self.in_channels = in_channels
self.out_channels = out_channels
self.normalize = normalize
self.root_weight = root_weight
self.project = project
if isinstance(in_channels, int):
in_channels = (in_channels, in_channels)
if self.project:
self.lin = Linear(in_channels[0], in_channels[0], bias=True)
if self.aggr is None:
self.fuse = False # No "fused" message_and_aggregate.
self.lstm = LSTM(in_channels[0], in_channels[0], batch_first=True)
self.lin_l = Linear(in_channels[0], out_channels, bias=bias)
if self.root_weight:
self.lin_r = Linear(in_channels[1], out_channels, bias=False)
self.reset_parameters()
def reset_parameters(self):
if self.project:
self.lin.reset_parameters()
if self.aggr is None:
self.lstm.reset_parameters()
self.lin_l.reset_parameters()
if self.root_weight:
self.lin_r.reset_parameters()
def forward(self,
x: Union[Tensor, OptPairTensor],
edge_index: Adj,
size: Size = None) -> Tensor:
""""""
if isinstance(x, Tensor):
x: OptPairTensor = (x, x)
if self.project and hasattr(self, 'lin'):
x = (self.lin(x[0]).relu(), x[1])
# propagate_type: (x: OptPairTensor)
out = self.propagate(edge_index, x=x, size=size)
out = self.lin_l(out)
x_r = x[1]
if self.root_weight and x_r is not None:
out += self.lin_r(x_r)
if self.normalize:
out = F.normalize(out, p=2., dim=-1)
return out
def message(self, x_j: Tensor) -> Tensor:
return x_j
def message_and_aggregate(self, adj_t: SparseTensor,
x: OptPairTensor) -> Tensor:
adj_t = adj_t.set_value(None, layout=None)
return matmul(adj_t, x[0], reduce=self.aggr)
def aggregate(self,
x: Tensor,
index: Tensor,
ptr: Optional[Tensor] = None,
dim_size: Optional[int] = None) -> Tensor:
if self.aggr is not None:
return scatter(x,
index,
dim=self.node_dim,
dim_size=dim_size,
reduce=self.aggr)
# LSTM aggregation:
if ptr is None and not torch.all(index[:-1] <= index[1:]):
raise ValueError(f"Can not utilize LSTM-style aggregation inside "
f"'{self.__class__.__name__}' in case the "
f"'edge_index' tensor is not sorted by columns. "
f"Run 'sort_edge_index(..., sort_by_row=False)' "
f"in a pre-processing step.")
x, mask = to_dense_batch(x, batch=index, batch_size=dim_size)
out, _ = self.lstm(x)
return out[:, -1]
def __repr__(self) -> str:
aggr = self.aggr if self.aggr is not None else 'lstm'
return (f'{self.__class__.__name__}({self.in_channels}, '
f'{self.out_channels}, aggr={aggr})')
