def __init__(self,
out_dim,
hidden_dim=16,
n_layers=4,
max_degree=6,
n_atom_types=MAX_ATOMIC_NUM,
concat_hidden=False):
super(NFP, self).__init__()
num_degree_type = max_degree + 1
with self.init_scope():
self.embed = EmbedAtomID(in_size=n_atom_types, out_size=hidden_dim)
self.layers = chainer.ChainList(*[
NFPUpdate(hidden_dim, hidden_dim, max_degree=max_degree)
for _ in range(n_layers)
])
self.read_out_layers = chainer.ChainList(
*[NFPReadout(hidden_dim, out_dim) for _ in range(n_layers)])
self.out_dim = out_dim
self.hidden_dim = hidden_dim
self.max_degree = max_degree
self.num_degree_type = num_degree_type
self.n_layers = n_layers
self.concat_hidden = concat_hidden
def model_with_nfp_no_dropout():
return RSGCN(out_dim=out_dim,
readout=NFPReadout(in_channels=out_dim, out_size=out_dim),
dropout_ratio=0.)
def model_with_nfp():
return RSGCN(out_dim=out_dim,
readout=NFPReadout(in_channels=out_dim, out_size=out_dim))
def readout():
return NFPReadout(in_channels=hidden_dim, out_size=out_dim)