Enable to use GGNNReadout for RelGCN
@@ -16,29 +16,38 @@ class GGNNReadout(chainer.Chain):
each layer and the result is concatenated
nobias (bool): If ``True``, then this function does not use
the bias
+ activate (~chainer.Function or ~chainer.FunctionNode
+ or NoneType):
+ You can use this class for RelGCN with functions.tanh
"""
def __init__(self, out_dim, hidden_dim=16, n_layers=4,
- concat_hidden=False):
+ concat_hidden=False, nobias=False, activate=None):
super(GGNNReadout, self).__init__()
n_layer = n_layers if concat_hidden else 1
with self.init_scope():
self.i_layers = chainer.ChainList(
- *[GraphLinear(2 * hidden_dim, out_dim)
+ *[GraphLinear(None, out_dim, nobias=nobias)
for _ in range(n_layer)])
self.j_layers = chainer.ChainList(
- *[GraphLinear(hidden_dim, out_dim) for _ in range(n_layer)])
+ *[GraphLinear(None, out_dim, nobias=nobias)
+ for _ in range(n_layer)])
self.out_dim = out_dim
self.hidden_dim = hidden_dim
self.n_layers = n_layers
self.concat_hidden = concat_hidden
+ self.nobias = nobias
+ self.activate = functions.identity if activate is None else activate
- def __call__(self, h, h0, step=0):
+ def __call__(self, h, h0=None, step=0):
# --- Readout part ---
index = step if self.concat_hidden else 0
# h, h0: (minibatch, atom, ch)
- g = functions.sigmoid(
- self.i_layers[index](functions.concat((h, h0), axis=2))) \
- * self.j_layers[index](h)
- g = functions.sum(g, axis=1) # sum along atom's axis
+ h1 = functions.concat((h, h0), axis=2) if h0 is not None else h
+
+ g1 = functions.sigmoid(self.i_layers[index](h1))
+ h2 = h if self.activate == functions.identity else h1
+ g2 = self.activate(self.j_layers[index](h2))
+ # sum along atom's axis
+ g = self.activate(functions.sum(g1 * g2, axis=1))
return g
