use GraphConvModel for simple relgat model
@@ -42,8 +42,9 @@ def rescale_adj(adj):
class GraphConvModel(chainer.Chain):
def __init__(self, in_channels, out_dim, update_layer, readout_layer, n_layers=None,
hidden_dim_super=None, n_atom_types=MAX_ATOMIC_NUM, n_edge_types=4, max_degree=6,
- dropout_ratio=-1.0, with_gwm=True, concat_hidden=False, sum_hidden=False,
- weight_tying=False, scale_adj=False, activation=None, use_batchnorm=False):
+ n_heads=8, negative_slope=0.2, dropout_ratio=-1.0, with_gwm=True,
+ concat_hidden=False, sum_hidden=False, weight_tying=False, scale_adj=False, activation=None,
+ use_batchnorm=False):
# Note: in_channels can be integer or list
# Note: Is out_dim necessary?
super(GraphConvModel, self).__init__()
@@ -74,7 +75,8 @@ class GraphConvModel(chainer.Chain):
self.embed = EmbedAtomID(out_size=in_channels[0], in_size=n_atom_types)
self.update_layers = chainer.ChainList(
*[update_layer(in_channels=in_channels[i], out_channels=in_channels[i+1],
- n_edge_types=n_edge_types, dropout_ratio=dropout_ratio)
+ n_edge_types=n_edge_types, dropout_ratio=dropout_ratio,
+ n_heads=n_heads, negative_slope=negative_slope)
for i in range(n_update_layers)])
self.readout_layers = chainer.ChainList(
*[readout_layer(out_dim=out_dim, in_channels=in_channels[-1])
