wehub-resource-sync
258 lines
8.4 KiB
Python
258 lines
8.4 KiB
Python
"""Torch module for grouped reversible residual connections for GNNs"""
|
|
# pylint: disable= no-member, arguments-differ, invalid-name, C0116, R1728
|
|
from copy import deepcopy
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class InvertibleCheckpoint(torch.autograd.Function):
|
|
r"""Extension of torch.autograd"""
|
|
|
|
@staticmethod
|
|
def forward(ctx, fn, fn_inverse, num_inputs, *inputs_and_weights):
|
|
ctx.fn = fn
|
|
ctx.fn_inverse = fn_inverse
|
|
ctx.weights = inputs_and_weights[num_inputs:]
|
|
inputs = inputs_and_weights[:num_inputs]
|
|
ctx.input_requires_grad = []
|
|
|
|
with torch.no_grad():
|
|
# Make a detached copy, which shares the storage
|
|
x = []
|
|
for element in inputs:
|
|
if isinstance(element, torch.Tensor):
|
|
x.append(element.detach())
|
|
ctx.input_requires_grad.append(element.requires_grad)
|
|
else:
|
|
x.append(element)
|
|
ctx.input_requires_grad.append(None)
|
|
# Detach the output, which then allows discarding the intermediary results
|
|
outputs = ctx.fn(*x).detach_()
|
|
|
|
# clear memory of input node features
|
|
inputs[1].untyped_storage().resize_(0)
|
|
|
|
# store for backward pass
|
|
ctx.inputs = [inputs]
|
|
ctx.outputs = [outputs]
|
|
|
|
return outputs
|
|
|
|
@staticmethod
|
|
def backward(ctx, *grad_outputs):
|
|
if not torch.autograd._is_checkpoint_valid():
|
|
raise RuntimeError(
|
|
"InvertibleCheckpoint is not compatible with .grad(), \
|
|
please use .backward() if possible"
|
|
)
|
|
# retrieve input and output tensor nodes
|
|
if len(ctx.outputs) == 0:
|
|
raise RuntimeError(
|
|
"Trying to perform backward on the InvertibleCheckpoint \
|
|
for more than once."
|
|
)
|
|
inputs = ctx.inputs.pop()
|
|
outputs = ctx.outputs.pop()
|
|
|
|
# reconstruct input node features
|
|
with torch.no_grad():
|
|
# inputs[0] is DGLGraph and inputs[1] is input node features
|
|
inputs_inverted = ctx.fn_inverse(
|
|
*((inputs[0], outputs) + inputs[2:])
|
|
)
|
|
# clear memory of outputs
|
|
outputs.untyped_storage().resize_(0)
|
|
|
|
x = inputs[1]
|
|
x.untyped_storage().resize_(int(np.prod(x.size())))
|
|
x.set_(inputs_inverted)
|
|
|
|
# compute gradients
|
|
with torch.set_grad_enabled(True):
|
|
detached_inputs = []
|
|
for i, element in enumerate(inputs):
|
|
if isinstance(element, torch.Tensor):
|
|
element = element.detach()
|
|
element.requires_grad = ctx.input_requires_grad[i]
|
|
detached_inputs.append(element)
|
|
|
|
detached_inputs = tuple(detached_inputs)
|
|
temp_output = ctx.fn(*detached_inputs)
|
|
|
|
filtered_detached_inputs = tuple(
|
|
filter(
|
|
lambda x: getattr(x, "requires_grad", False), detached_inputs
|
|
)
|
|
)
|
|
gradients = torch.autograd.grad(
|
|
outputs=(temp_output,),
|
|
inputs=filtered_detached_inputs + ctx.weights,
|
|
grad_outputs=grad_outputs,
|
|
)
|
|
|
|
input_gradients = []
|
|
i = 0
|
|
for rg in ctx.input_requires_grad:
|
|
if rg:
|
|
input_gradients.append(gradients[i])
|
|
i += 1
|
|
else:
|
|
input_gradients.append(None)
|
|
|
|
gradients = tuple(input_gradients) + gradients[-len(ctx.weights) :]
|
|
|
|
return (None, None, None) + gradients
|
|
|
|
|
|
class GroupRevRes(nn.Module):
|
|
r"""Grouped reversible residual connections for GNNs, as introduced in
|
|
`Training Graph Neural Networks with 1000 Layers <https://arxiv.org/abs/2106.07476>`__
|
|
|
|
It uniformly partitions an input node feature :math:`X` into :math:`C` groups
|
|
:math:`X_1, X_2, \cdots, X_C` across the channel dimension. Besides, it makes
|
|
:math:`C` copies of the input GNN module :math:`f_{w1}, \cdots, f_{wC}`. In the
|
|
forward pass, each GNN module only takes the corresponding group of node features.
|
|
|
|
The output node representations :math:`X^{'}` are computed as follows.
|
|
|
|
.. math::
|
|
|
|
X_0^{'} = \sum_{i=2}^{C}X_i
|
|
|
|
X_i^{'} = f_{wi}(X_{i-1}^{'}, g, U) + X_i, i\in\{1,\cdots,C\}
|
|
|
|
X^{'} = X_1^{'} \, \Vert \, \ldots \, \Vert \, X_C^{'}
|
|
|
|
where :math:`g` is the input graph, :math:`U` is arbitrary additional input arguments like
|
|
edge features, and :math:`\, \Vert \,` is concatenation.
|
|
|
|
Parameters
|
|
----------
|
|
gnn_module : nn.Module
|
|
GNN module for message passing. :attr:`GroupRevRes` will clone the module for
|
|
:attr:`groups`-1 number of times, yielding :attr:`groups` copies in total.
|
|
The input and output node representation size need to be the same. Its forward
|
|
function needs to take a DGLGraph and the associated input node features in order,
|
|
optionally followed by additional arguments like edge features.
|
|
groups : int, optional
|
|
The number of groups.
|
|
|
|
Examples
|
|
--------
|
|
|
|
>>> import dgl
|
|
>>> import torch
|
|
>>> import torch.nn as nn
|
|
>>> from dgl.nn import GraphConv, GroupRevRes
|
|
|
|
>>> class GNNLayer(nn.Module):
|
|
... def __init__(self, feats, dropout=0.2):
|
|
... super(GNNLayer, self).__init__()
|
|
... # Use BatchNorm and dropout to prevent gradient vanishing
|
|
... # In particular if you use a large number of GNN layers
|
|
... self.norm = nn.BatchNorm1d(feats)
|
|
... self.conv = GraphConv(feats, feats)
|
|
... self.dropout = nn.Dropout(dropout)
|
|
...
|
|
... def forward(self, g, x):
|
|
... x = self.norm(x)
|
|
... x = self.dropout(x)
|
|
... return self.conv(g, x)
|
|
|
|
>>> num_nodes = 5
|
|
>>> num_edges = 20
|
|
>>> feats = 32
|
|
>>> groups = 2
|
|
>>> g = dgl.rand_graph(num_nodes, num_edges)
|
|
>>> x = torch.randn(num_nodes, feats)
|
|
>>> conv = GNNLayer(feats // groups)
|
|
>>> model = GroupRevRes(conv, groups)
|
|
>>> out = model(g, x)
|
|
"""
|
|
|
|
def __init__(self, gnn_module, groups=2):
|
|
super(GroupRevRes, self).__init__()
|
|
self.gnn_modules = nn.ModuleList()
|
|
for i in range(groups):
|
|
if i == 0:
|
|
self.gnn_modules.append(gnn_module)
|
|
else:
|
|
self.gnn_modules.append(deepcopy(gnn_module))
|
|
self.groups = groups
|
|
|
|
def _forward(self, g, x, *args):
|
|
xs = torch.chunk(x, self.groups, dim=-1)
|
|
|
|
if len(args) == 0:
|
|
args_chunks = [()] * self.groups
|
|
else:
|
|
chunked_args = list(
|
|
map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
|
|
)
|
|
args_chunks = list(zip(*chunked_args))
|
|
y_in = sum(xs[1:])
|
|
|
|
ys = []
|
|
for i in range(self.groups):
|
|
y_in = xs[i] + self.gnn_modules[i](g, y_in, *args_chunks[i])
|
|
ys.append(y_in)
|
|
|
|
out = torch.cat(ys, dim=-1)
|
|
|
|
return out
|
|
|
|
def _inverse(self, g, y, *args):
|
|
ys = torch.chunk(y, self.groups, dim=-1)
|
|
|
|
if len(args) == 0:
|
|
args_chunks = [()] * self.groups
|
|
else:
|
|
chunked_args = list(
|
|
map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
|
|
)
|
|
args_chunks = list(zip(*chunked_args))
|
|
|
|
xs = []
|
|
for i in range(self.groups - 1, -1, -1):
|
|
if i != 0:
|
|
y_in = ys[i - 1]
|
|
else:
|
|
y_in = sum(xs)
|
|
|
|
x = ys[i] - self.gnn_modules[i](g, y_in, *args_chunks[i])
|
|
xs.append(x)
|
|
|
|
x = torch.cat(xs[::-1], dim=-1)
|
|
|
|
return x
|
|
|
|
def forward(self, g, x, *args):
|
|
r"""Apply the GNN module with grouped reversible residual connection.
|
|
|
|
Parameters
|
|
----------
|
|
g : DGLGraph
|
|
The graph.
|
|
x : torch.Tensor
|
|
The input feature of shape :math:`(N, D_{in})`, where :math:`D_{in}` is size
|
|
of input feature, :math:`N` is the number of nodes.
|
|
args
|
|
Additional arguments to pass to :attr:`gnn_module`.
|
|
|
|
Returns
|
|
-------
|
|
torch.Tensor
|
|
The output feature of shape :math:`(N, D_{in})`.
|
|
"""
|
|
args = (g, x) + args
|
|
y = InvertibleCheckpoint.apply(
|
|
self._forward,
|
|
self._inverse,
|
|
len(args),
|
|
*(args + tuple([p for p in self.parameters() if p.requires_grad]))
|
|
)
|
|
|
|
return y
|