def reaction_encoder(dilations, gf_dim, kernel, reactions, pooling, reuse,
scope_name, input_width):
compressed = []
strides_schedule = [(2, 2), (2, 2), (2, 2)]
down_layers = len(strides_schedule)
with tf.variable_scope("down", reuse=tf.AUTO_REUSE):
for i in range(4):
hidden_dim = gf_dim / 8
reaction_part = tf.transpose(ops.sn_embedding(
tf.squeeze(reactions[:, i, :]),
len(SMILES_CHARACTER_TO_ID),
8,
name='component_embedding'),
perm=[0, 2, 1])
tf.summary.histogram("embedded_reactions",
reaction_part,
family=scope_name)
h = tf.expand_dims(reaction_part, axis=3)
if input_width < 256:
h = tf.nn.leaky_relu(
ops.snconv2d(h,
hidden_dim, (h.get_shape().as_list()[1], 3),
name='first_component_conv',
padding="CYCLE",
strides=(1, 2)))
tf.summary.histogram("first_component_conv", h, family=scope_name)
for layer_id in range(down_layers):
hidden_dim = hidden_dim * strides_schedule[layer_id][0]
block_name = 'block{}'.format(layer_id)
dilation_rate = dilations[0]**(layer_id +
1), dilations[1]**(layer_id + 1)
h = ops.sn_norm_block(h,
hidden_dim,
block_name,
get_kernel(h, kernel),
strides_schedule[layer_id],
dilation_rate,
act=tf.nn.leaky_relu,
pooling=pooling,
padding="VALID")
tf.summary.histogram(block_name, h, family=scope_name)
if layer_id == 0:
h = attention(h, hidden_dim, sn=True, reuse=reuse)
tf.summary.histogram("attention", h, family=scope_name)
compressed.append(h)
print("Component Shape", compressed[0].shape)
compressed_reaction = tf.concat(
[compressed[0], compressed[1], compressed[2], compressed[3]],
axis=3,
name="compressed_reaction")
tf.summary.histogram("compressed_reaction",
compressed_reaction,
family=scope_name)
print("REACTION COMPRESSED TO: ", compressed_reaction.shape)
return compressed_reaction
def generator_resnet(zs,
labels,
gf_dim,
num_classes,
kernel=(3, 3),
strides=[(2, 2)],
dilations=(1, 1),
pooling='avg',
scope_name='Generator',
reuse=False,
output_shape=[8, 128, 1],
act=tf.nn.leaky_relu):
with tf.variable_scope(scope_name, reuse=tf.AUTO_REUSE) as scope:
height_d, width_d = get_dimentions_factors(strides)
number_of_layers = len(strides)
hidden_dim = gf_dim * (2**(number_of_layers - 1))
c_h = int(output_shape[0] / height_d)
c_w = int((output_shape[1] / width_d))
h = ops.snlinear(zs, c_h * c_w * hidden_dim, name='noise_linear')
h = tf.reshape(h, [-1, c_h, c_w, hidden_dim])
print("COMPRESSED TO: ", h.shape)
with tf.variable_scope("up", reuse=tf.AUTO_REUSE):
for layer_id in range(number_of_layers):
print(h.shape)
block_name = 'up_block{}'.format(number_of_layers -
(layer_id + 1))
dilation_rate = (dilations[0]**(number_of_layers - layer_id),
dilations[1]**(number_of_layers - layer_id))
h = sn_block(h, hidden_dim, block_name, get_kernel(h, kernel),
strides[layer_id], dilation_rate, act, pooling,
'VALID')
tf.summary.histogram(block_name, h, family=scope_name)
if layer_id == number_of_layers - 2:
h = attention(h, hidden_dim, sn=True, reuse=reuse)
tf.summary.histogram("up_attention", h, family=scope_name)
hidden_dim = hidden_dim / strides[layer_id][1]
bn = ops.BatchNorm(name='g_bn')
h_act = tf.nn.leaky_relu(bn(h), name="h_act")
if output_shape[2] == 1:
out = tf.nn.tanh(ops.snconv2d(h_act,
1, (output_shape[0], 1),
name='last_conv'),
name="generated")
else:
out = tf.nn.tanh(ops.snconv2d(h_act, 3, (1, 1), name='last_conv'),
name="generated")
tf.summary.histogram("Generated_results", out, family=scope_name)
print("GENERATED SHAPE", out.shape)
return out
def discriminator_resnet(x,
labels,
df_dim,
number_classes,
kernel=(3, 3),
strides=[(2, 2)],
dilations=(1, 1),
pooling='avg',
update_collection=None,
act=tf.nn.relu,
scope_name='Discriminator',
reuse=False):
with tf.variable_scope(scope_name) as scope:
if reuse:
scope.reuse_variables()
tf.summary.histogram("Input", x, family=scope_name)
h = x
hidden_dim = df_dim
for layer in range(len(strides)):
print(h.shape)
if layer == 1:
h = attention(h, hidden_dim, sn=True, reuse=reuse)
tf.summary.histogram("attention", h, family=scope_name)
block_name = 'd_block{}'.format(layer)
hidden_dim = hidden_dim * strides[layer][0]
dilation_rate = dilations[0]**(layer + 1), dilations[1]**(layer +
1)
h = sn_block(h,
hidden_dim,
block_name,
get_kernel(h, kernel),
strides[layer],
dilation_rate,
update_collection,
act,
pooling,
padding='VALID')
tf.summary.histogram(block_name, h, family=scope_name)
end_block = act(h, name="after_resnet_block")
tf.summary.histogram("after_resnet_block",
end_block,
family=scope_name)
# h_std = ops.minibatch_stddev_layer(end_block)
# tf.summary.histogram(h_std.name, h_std, family=scope_name)
# h_std_conv_std = act(ops.snconv2d(h_std, hidden_dim, (1, 3), update_collection=update_collection,
# name='minibatch_stddev_stride', padding=None, strides=(1, 3)),
# name="minibatch_stddev_stride_act")
# tf.summary.histogram("after_mini_batch_std", h_std_conv_std, family=scope_name)
# h_final_flattened = tf.layers.flatten(h_std_conv_std)
h_final_flattened = tf.reduce_sum(end_block, [1, 2])
tf.summary.histogram("h_final_flattened",
h_final_flattened,
family=scope_name)
output = ops.snlinear(h_final_flattened,
1,
update_collection=update_collection,
name='d_sn_linear')
tf.summary.histogram("final_output", output, family=scope_name)
# output = tf.Print(output, [tf.py_func(
# lambda val, score: print_protein_values(val,score),
# [tf.squeeze(x)[0], output[0]], tf.string)], "seq:")
return output, h_final_flattened
def generator_1d(zs,
labels,
gf_dim,
num_classes,
kernel=(3, 3),
strides=(2, 2),
dilations=(1, 1),
pooling='avg',
scope_name='Generator',
reuse=False,
output_shape=[8, 128, 1],
act=tf.nn.leaky_relu):
with tf.variable_scope(scope_name, reuse=tf.AUTO_REUSE) as scope:
# strides_schedule = get_strides_schedule(input_width)
strides_schedule = [(1, 2), (1, 2), (1, 2), (1, 2)]
height_d, width_d = get_dimentions_factors(strides_schedule)
number_of_layers = len(strides_schedule)
hidden_dim = gf_dim * (2**(number_of_layers - 1))
kernel = (1, 5)
c_h = 1 # int(embedding_height / height_d)
c_w = int((output_shape[1] / width_d))
h = ops.snlinear(zs, c_h * int(c_w) * hidden_dim, name='noise_linear')
tf.summary.histogram("noise_snlinear", h, family=scope_name)
h = tf.reshape(h, shape=[-1, c_h, int(c_w), hidden_dim])
# new_shape = [h.get_shape().as_list()[0], h.get_shape().as_list()[1],
# h.get_shape().as_list()[2] * 2, hidden_dim]
# h = ops.sndeconv2d(h, new_shape, (c_h, 3), name='noise_expand', strides=(1, 2))
# h = act(h)
tf.summary.histogram("before_blocks", h, family=scope_name)
print("COMPRESSED TO: ", h.shape)
with tf.variable_scope("up", reuse=tf.AUTO_REUSE):
for layer_id in range(number_of_layers):
block_name = 'up_block{}'.format(number_of_layers -
(layer_id + 1))
# dilation_rate = (dilations[0] ** (number_of_layers - layer_id),
# dilations[1] ** (number_of_layers - layer_id))
dilation_rate = (1, 1)
h = sn_block(h, hidden_dim, block_name, kernel,
strides_schedule[layer_id], dilation_rate, act,
pooling, 'VALID')
tf.summary.histogram(block_name, h, family=scope_name)
if layer_id == number_of_layers - 2:
h = attention(h, hidden_dim, sn=True, reuse=reuse)
tf.summary.histogram("up_attention", h, family=scope_name)
hidden_dim = hidden_dim / 2
bn = ops.BatchNorm(name='g_bn')
h_act = tf.nn.leaky_relu(bn(h), name="h_act")
h_act = tf.reshape(h_act, [
-1, output_shape[0], output_shape[1],
int(hidden_dim / output_shape[0]) * 2
])
h_act = ops.snconv2d(h_act,
int(hidden_dim / output_shape[0]) * 2,
(output_shape[0], 1),
name='embedding_height')
out = tf.nn.tanh(ops.snconv2d(h_act,
1, (output_shape[0], 1),
name='last_conv'),
name="generated")
tf.summary.histogram("Generated_results", out, family=scope_name)
print("GENERATED SHAPE", out.shape)
return out