def _build_graph_conv(self, F_h):
self.F_h = list(F_h) if isinstance(F_h, tuple) else F_h
self.conv, self.bn = [], []
for i, (f_in,
f_out) in enumerate(zip([self.F_e] + self.F_h[:-1], self.F_h)):
conv = modules.GraphConv(f_in, f_out, self.N_B + self.D)
self.conv.append(conv)
self.register_child(conv)
if i != 0:
bn = modules.BatchNorm(in_channels=f_in)
self.register_child(bn)
else:
bn = None
self.bn.append(bn)
self.bn_skip = modules.BatchNorm(in_channels=sum(self.F_h))
self.linear_skip = modules.Linear_BN(sum(self.F_h), self.F_skip)
# projectors for conditional variable
self.linear_c = []
for i, f_out in enumerate(self.F_h):
if self.rename:
linear_c = nn.Dense(f_out,
use_bias=False,
in_units=self.N_C,
prefix='cond_{}'.format(i))
else:
linear_c = nn.Dense(f_out, use_bias=False, in_units=self.N_C)
self.register_child(linear_c)
self.linear_c.append(linear_c)
def _build_graph_conv(self, F_h):
self.F_h = list(F_h) if isinstance(F_h, tuple) else F_h
self.conv, self.bn = [], []
for i, (f_in,
f_out) in enumerate(zip([self.F_e] + self.F_h[:-1], self.F_h)):
conv = modules.GraphConv(f_in, f_out, self.N_B + self.D)
self.conv.append(conv)
self.register_child(conv)
if i != 0:
bn = modules.BatchNorm(in_channels=f_in)
self.register_child(bn)
else:
bn = None
self.bn.append(bn)
self.bn_skip = modules.BatchNorm(in_channels=sum(self.F_h))
self.linear_skip = modules.Linear_BN(sum(self.F_h), self.F_skip)
def __init__(self, N_A, N_B, D, F_e, F_skip, F_c, Fh_policy, activation,
*args, **kwargs):
super(MoleculeGenerator, self).__init__()
self.N_A = N_A
self.N_B = N_B
self.D = D
self.F_e = F_e
self.F_skip = F_skip
self.F_c = list(F_c) if isinstance(F_c, tuple) else F_c
self.Fh_policy = Fh_policy
self.activation = fn.get_activation(activation)
with self.name_scope():
# embeddings
self.embedding_atom = nn.Embedding(self.N_A, self.F_e)
self.embedding_mask = nn.Embedding(3, self.F_e)
# graph conv
self._build_graph_conv(*args, **kwargs)
# fully connected
self.dense = nn.Sequential()
for i, (f_in, f_out) in enumerate(
zip([
self.F_skip,
] + self.F_c[:-1], self.F_c)):
self.dense.add(modules.Linear_BN(f_in, f_out))
# policy
self.policy_0 = self.params.get('policy_0',
shape=[
self.N_A,
],
init=mx.init.Zero(),
allow_deferred_init=False)
self.policy_h = modules.Policy(self.F_c[-1], self.Fh_policy,
self.N_A, self.N_B)
self.mode = 'loss'
