File size: 1,521 Bytes
26225c5 |
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 |
from torch_scatter import scatter
from torch import nn
from src.data import NAG
from src.nn import MLP, BatchNorm
__all__ = ['NodeMLP']
class NodeMLP(nn.Module):
"""Simple MLP on the handcrafted features of the level-i in a NAG.
This is used as a baseline to test how expressive handcrafted
features are.
"""
def __init__(
self, dims, level=0, activation=nn.LeakyReLU(), norm=BatchNorm,
drop=None, norm_mode='graph'):
super().__init__()
self.level = level
self.mlp = MLP(dims, activation=activation, norm=norm, drop=drop)
self.norm_mode = norm_mode
@property
def out_dim(self):
return self.mlp.out_dim
def forward(self, nag):
assert isinstance(nag, NAG)
assert nag.num_levels > self.level
# Compute node features from the handcrafted features
norm_index = nag[self.i_level].norm_index(mode=self.norm_mode)
x = self.mlp(nag[self.level].x, batch=norm_index)
# If node level is 1, output level-1 features
if self.level == 1:
return x
# If node level is 0, max-pool to produce level-1 features
if self.level == 0:
return scatter(
x, nag[0].super_index, dim=0, dim_size=nag[1].num_nodes,
reduce='max')
# If node level is larger than 1, distribute parent features to
# level-1 nodes
super_index = nag.get_super_index(self.level, low=1)
return x[super_index]
|