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)