wehub-resource-sync
56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
"""
|
|
Recurrent Relational Network(RRN) module
|
|
|
|
References:
|
|
- Recurrent Relational Networks
|
|
- Paper: https://arxiv.org/abs/1711.08028
|
|
- Original Code: https://github.com/rasmusbergpalm/recurrent-relational-networks
|
|
"""
|
|
|
|
import dgl.function as fn
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class RRNLayer(nn.Module):
|
|
def __init__(self, msg_layer, node_update_func, edge_drop):
|
|
super(RRNLayer, self).__init__()
|
|
self.msg_layer = msg_layer
|
|
self.node_update_func = node_update_func
|
|
self.edge_dropout = nn.Dropout(edge_drop)
|
|
|
|
def forward(self, g):
|
|
g.apply_edges(self.get_msg)
|
|
g.edata["e"] = self.edge_dropout(g.edata["e"])
|
|
g.update_all(
|
|
message_func=fn.copy_e("e", "msg"), reduce_func=fn.sum("msg", "m")
|
|
)
|
|
g.apply_nodes(self.node_update)
|
|
|
|
def get_msg(self, edges):
|
|
e = torch.cat([edges.src["h"], edges.dst["h"]], -1)
|
|
e = self.msg_layer(e)
|
|
return {"e": e}
|
|
|
|
def node_update(self, nodes):
|
|
return self.node_update_func(nodes)
|
|
|
|
|
|
class RRN(nn.Module):
|
|
def __init__(self, msg_layer, node_update_func, num_steps, edge_drop):
|
|
super(RRN, self).__init__()
|
|
self.num_steps = num_steps
|
|
self.rrn_layer = RRNLayer(msg_layer, node_update_func, edge_drop)
|
|
|
|
def forward(self, g, get_all_outputs=True):
|
|
outputs = []
|
|
for _ in range(self.num_steps):
|
|
self.rrn_layer(g)
|
|
if get_all_outputs:
|
|
outputs.append(g.ndata["h"])
|
|
if get_all_outputs:
|
|
outputs = torch.stack(outputs, 0) # num_steps x n_nodes x h_dim
|
|
else:
|
|
outputs = g.ndata["h"] # n_nodes x h_dim
|
|
return outputs
|