Spaces:
Runtime error
Runtime error
File size: 2,442 Bytes
d2a8669 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import math
import torch
from torch import Tensor
from torch.nn import Parameter
from torch_geometric.nn.conv import MessagePassing, GCNConv
class GCNIIConv(MessagePassing):
"""
The graph convolutional operator with initial residual connections and
identity mapping (GCNII) from the `"Simple and Deep Graph Convolutional
Networks" <https://arxiv.org/abs/2007.02133>`_ paper
"""
def __init__(self, channels, alpha, theta=None, layer=None,
shared_weights=True, cached=False, **kwargs):
super(GCNIIConv, self).__init__(aggr='add', **kwargs)
self.channels = channels
self.alpha = alpha
self.beta = 1.
if theta is not None or layer is not None:
assert theta is not None and layer is not None
self.beta = math.log(theta / layer + 1)
# self.cached = cached
self.weight1 = Parameter(torch.Tensor(channels, channels))
if shared_weights:
self.register_parameter('weight2', None)
else:
self.weight2 = Parameter(torch.Tensor(channels, channels))
self.reset_parameters()
def reset_parameters(self):
self.glorot(self.weight1)
self.glorot(self.weight2)
def glorot(self, tensor):
if tensor is not None:
stdv = math.sqrt(6.0 / (tensor.size(-2) + tensor.size(-1)))
tensor.data.uniform_(-stdv, stdv)
def forward(self, x, x_0, edge_index, edge_weight=None):
edge_index, norm = GCNConv.norm(edge_index, x.size(0), edge_weight,
dtype=x.dtype)
x = self.propagate(edge_index, x=x, norm=norm)
if self.weight2 is None:
out = (1 - self.alpha) * x + self.alpha * x_0
out = (1 - self.beta) * out + self.beta * (out @ self.weight1)
else:
out1 = (1 - self.alpha) * x
out1 = (1 - self.beta) * out1 + self.beta * (out1 @ self.weight1)
out2 = self.alpha * x_0
out2 = (1 - self.beta) * out2 + self.beta * (out2 @ self.weight2)
out = out1 + out2
return out
def message(self, x_j, norm):
return norm.view(-1, 1) * x_j
def __repr__(self):
return '{}({}, alpha={}, beta={})'.format(self.__class__.__name__,
self.channels, self.alpha,
self.beta) |