chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,257 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user