def call(self, inputs): data_a, data_x, data = inputs self.embeddings = self.sample_z(self.batch) # generate self.edges_logits, self.nodes_logits = self.decoder( self.embeddings, self.decoder_units, self.vertexes, self.edges, self.nodes, training=False, dropout_rate=self.dropout_rate) (self.edges_softmax, self.nodes_softmax), \ (self.edges_argmax, self.nodes_argmax), \ (self.edges_gumbel_logits, self.nodes_gumbel_logits), \ (self.edges_gumbel_softmax, self.nodes_gumbel_softmax), \ (self.edges_gumbel_argmax, self.nodes_gumbel_argmax) = postprocess_logits( (self.edges_logits, self.nodes_logits), temperature=self.temperature) self.adjacency_tensor = tf.one_hot(data_a, depth=self.edges, dtype=tf.float32) self.node_tensor = tf.one_hot(data_x, depth=self.nodes, dtype=tf.float32) # self.rewardR = reward(mols) # 采样分子 n, e = self.nodes_gumbel_argmax, self.edges_gumbel_argmax n, e = np.argmax(n, axis=-1), np.argmax(e, axis=-1) gen_mols = [ data.matrices2mol(n_, e_, strict=True) for n_, e_ in zip(n, e) ] self.rewardF = reward(gen_mols) # 生成分子 self.edges_hat = tf.case( [ (self.soft_gumbel_softmax, lambda: self.edges_gumbel_softmax), # (self.hard_gumbel_softmax, lambda: tf.stop_gradient( self.edges_gumbel_argmax - self.edges_gumbel_softmax) + self.edges_gumbel_softmax) ], # default=lambda: self.edges_softmax, exclusive=True) self.nodes_hat = tf.case( [(self.soft_gumbel_softmax, lambda: self.nodes_gumbel_softmax), (self.hard_gumbel_softmax, lambda: tf.stop_gradient( self.nodes_gumbel_argmax - self.nodes_gumbel_softmax) + self. nodes_gumbel_softmax)], default=lambda: self.nodes_softmax, exclusive=True) self.logits_real, self.features_real = self.D_x( (self.adjacency_tensor, None, self.node_tensor), units=self.discriminator_units) self.logits_fake, self.features_fake = self.D_x( (self.edges_hat, None, self.nodes_hat), units=self.discriminator_units) self.value_logits_real = self.V_x( (self.adjacency_tensor, None, self.node_tensor), units=self.discriminator_units) self.value_logits_fake = self.V_x( (self.edges_hat, None, self.nodes_hat), units=self.discriminator_units)
def __init__(self, vertexes, edges, nodes, embedding_dim, decoder_units, discriminator_units, decoder, discriminator, soft_gumbel_softmax=False, hard_gumbel_softmax=False, batch_discriminator=True): self.vertexes, self.edges, self.nodes, self.embedding_dim, self.decoder_units, self.discriminator_units, \ self.decoder, self.discriminator, self.batch_discriminator = vertexes, edges, nodes, embedding_dim, decoder_units, \ discriminator_units, decoder, discriminator, batch_discriminator self.training = tf.placeholder_with_default(False, shape=()) self.dropout_rate = tf.placeholder_with_default(0., shape=()) self.soft_gumbel_softmax = tf.placeholder_with_default( soft_gumbel_softmax, shape=()) self.hard_gumbel_softmax = tf.placeholder_with_default( hard_gumbel_softmax, shape=()) self.temperature = tf.placeholder_with_default(1., shape=()) self.edges_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes, vertexes)) self.nodes_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes)) self.embeddings = tf.placeholder(dtype=tf.float32, shape=(None, embedding_dim)) #self.rewardR = tf.placeholder(dtype=tf.float32, shape=(None, 1)) #self.rewardF = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.adjacency_tensor = tf.one_hot(self.edges_labels, depth=edges, dtype=tf.float32) self.node_tensor = tf.one_hot(self.nodes_labels, depth=nodes, dtype=tf.float32) with tf.variable_scope('generator'): self.edges_logits, self.nodes_logits = self.decoder( self.embeddings, decoder_units, vertexes, edges, nodes, training=self.training, dropout_rate=self.dropout_rate) with tf.name_scope('outputs'): (self.edges_softmax, self.nodes_softmax), \ (self.edges_argmax, self.nodes_argmax), \ (self.edges_gumbel_logits, self.nodes_gumbel_logits), \ (self.edges_gumbel_softmax, self.nodes_gumbel_softmax), \ (self.edges_gumbel_argmax, self.nodes_gumbel_argmax) = postprocess_logits( (self.edges_logits, self.nodes_logits), temperature=self.temperature) self.edges_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.edges_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.edges_gumbel_argmax - self. edges_gumbel_softmax) + self. edges_gumbel_softmax }, default=lambda: self.edges_softmax, exclusive=True) self.nodes_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.nodes_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.nodes_gumbel_argmax - self. nodes_gumbel_softmax) + self. nodes_gumbel_softmax }, default=lambda: self.nodes_softmax, exclusive=True) with tf.name_scope('D_x_real'): self.logits_real, self.features_real = self.D_x( (self.adjacency_tensor, None, self.node_tensor), units=discriminator_units) with tf.name_scope('D_x_fake'): self.logits_fake, self.features_fake = self.D_x( (self.edges_hat, None, self.nodes_hat), units=discriminator_units)
def __init__(self, vertexes, edges, nodes, features, embedding_dim, encoder_units, decoder_units, variational, encoder, decoder, soft_gumbel_softmax=False, hard_gumbel_softmax=False, with_features=True): """ :param vertexes: the atoms num of molecular :param edges: the bond num :param nodes: the atom num type :param features: :param embedding_dim: :param encoder_units: :param decoder_units: :param variational: :param encoder: :param decoder: :param soft_gumbel_softmax: :param hard_gumbel_softmax: :param with_features: """ self.vertexes, self.nodes, self.edges, self.embedding_dim, self.encoder, self.decoder = \ vertexes, nodes, edges, embedding_dim, encoder, decoder self.training = tf.placeholder_with_default(False, shape=()) self.variational = tf.placeholder_with_default(variational, shape=()) self.soft_gumbel_softmax = tf.placeholder_with_default( soft_gumbel_softmax, shape=()) self.hard_gumbel_softmax = tf.placeholder_with_default( hard_gumbel_softmax, shape=()) self.temperature = tf.placeholder_with_default(1., shape=()) self.edges_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes, vertexes)) self.nodes_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes)) self.node_features = tf.placeholder(dtype=tf.float32, shape=(None, vertexes, features)) self.rewardR = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.rewardF = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.adjacency_tensor = tf.one_hot(self.edges_labels, depth=edges, dtype=tf.float32) self.node_tensor = tf.one_hot(self.nodes_labels, depth=nodes, dtype=tf.float32) with tf.variable_scope('encoder'): outputs = self.encoder( (self.adjacency_tensor, self.node_features if with_features else None, self.node_tensor), units=encoder_units[:-1], training=self.training, dropout_rate=0.) outputs = multi_dense_layers(outputs, units=encoder_units[-1], activation=tf.nn.tanh, training=self.training, dropout_rate=0.) self.embeddings_mean = tf.layers.dense(outputs, embedding_dim, activation=None) self.embeddings_std = tf.layers.dense(outputs, embedding_dim, activation=tf.nn.softplus) self.q_z = tf.distributions.Normal(self.embeddings_mean, self.embeddings_std) self.embeddings = tf.cond(self.variational, lambda: self.q_z.sample(), lambda: self.embeddings_mean) with tf.variable_scope('decoder'): self.edges_logits, self.nodes_logits = self.decoder( self.embeddings, decoder_units, vertexes, edges, nodes, training=self.training, dropout_rate=0.) with tf.name_scope('outputs'): (self.edges_softmax, self.nodes_softmax), \ (self.edges_argmax, self.nodes_argmax), \ (self.edges_gumbel_logits, self.nodes_gumbel_logits), \ (self.edges_gumbel_softmax, self.nodes_gumbel_softmax), \ (self.edges_gumbel_argmax, self.nodes_gumbel_argmax) = postprocess_logits( (self.edges_logits, self.nodes_logits), temperature=self.temperature) self.edges_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.edges_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.edges_gumbel_argmax - self. edges_gumbel_softmax) + self. edges_gumbel_softmax }, default=lambda: self.edges_softmax, exclusive=True) self.nodes_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.nodes_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.nodes_gumbel_argmax - self. nodes_gumbel_softmax) + self. nodes_gumbel_softmax }, default=lambda: self.nodes_softmax, exclusive=True) with tf.name_scope('V_x_real'): self.value_logits_real = self.V_x( (self.adjacency_tensor, None, self.node_tensor), units=encoder_units) with tf.name_scope('V_x_fake'): self.value_logits_fake = self.V_x( (self.edges_hat, None, self.nodes_hat), units=encoder_units)
def call(self, inputs=None, training=False): if inputs is not None: assert len( inputs ) == 2, "inputs must contain adjacent and node feature matrix." data_a, data_x = inputs # decoder_units, vertexes, edges, nodes = inputs sample_encoding = self.sample_z(self.batch) edges_logits, nodes_logits = self.G_x(sample_encoding) (edges_softmax, nodes_softmax), \ (edges_argmax, nodes_argmax), \ (edges_gumbel_logits, nodes_gumbel_logits), \ (edges_gumbel_softmax, nodes_gumbel_softmax), \ (edges_gumbel_argmax, nodes_gumbel_argmax) = postprocess_logits((edges_logits, nodes_logits), temperature=self.temperature) edges_hat = tf.case( [(self.soft_gumbel_softmax, lambda: edges_gumbel_softmax), (self.hard_gumbel_softmax, lambda: tf.stop_gradient( edges_gumbel_argmax - edges_gumbel_softmax) + edges_gumbel_softmax)], default=lambda: edges_softmax, exclusive=True) nodes_hat = tf.case( [(self.soft_gumbel_softmax, lambda: nodes_gumbel_softmax), (self.hard_gumbel_softmax, lambda: tf.stop_gradient( nodes_gumbel_argmax - nodes_gumbel_softmax) + nodes_gumbel_softmax)], default=lambda: nodes_softmax, exclusive=True) if not training: return nodes_gumbel_argmax, edges_gumbel_argmax # sample real if inputs is None: raise ValueError adjacency_tensor = tf.one_hot(data_a, depth=self.edges, dtype=tf.float32) node_tensor = tf.one_hot(data_x, depth=self.nodes, dtype=tf.float32) logits_real, features_real = self.D_x( (adjacency_tensor, None, node_tensor), training=training) value_logits_real, _ = self.V_x((adjacency_tensor, None, node_tensor), training=training) # generate logits_fake, features_fake = self.D_x((edges_hat, None, nodes_hat), training=training) value_logits_fake, _ = self.V_x((edges_hat, None, nodes_hat), training=training) # print("D_x: ", self.D_x.trainable_variables) real_sample = (adjacency_tensor, node_tensor, logits_real, features_real, value_logits_real) gen_sample = (edges_hat, nodes_hat, logits_fake, features_fake, value_logits_fake) gen_raw = (edges_softmax, nodes_softmax, edges_argmax, nodes_argmax, edges_gumbel_logits, nodes_gumbel_logits, edges_gumbel_softmax, nodes_gumbel_softmax, edges_gumbel_argmax, nodes_gumbel_argmax) return real_sample, gen_sample, gen_raw
def __init__( self, vertexes, # max atom num? edges, # bond num types nodes, # atom num types embedding_dim, # z space dim decoder_units, # tuple, decoder setup (z = Dense(z, dim=units_k)^{(k)}) discriminator_units, # tuple, discr. setup (GCN units, Readout units, MLP units) decoder, # callable fn. in models.__init__ discriminator, # callable fn. in models.__init__ soft_gumbel_softmax=False, hard_gumbel_softmax=False, batch_discriminator=True): self.vertexes, self.edges, self.nodes, \ self.embedding_dim, self.decoder_units, self.discriminator_units, \ self.decoder, self.discriminator, self.batch_discriminator = \ vertexes, edges, nodes, embedding_dim, decoder_units, \ discriminator_units, decoder, discriminator, batch_discriminator self.training = tf.placeholder_with_default(False, shape=()) self.dropout_rate = tf.placeholder_with_default(0., shape=()) self.soft_gumbel_softmax = tf.placeholder_with_default( soft_gumbel_softmax, shape=()) self.hard_gumbel_softmax = tf.placeholder_with_default( hard_gumbel_softmax, shape=()) self.temperature = tf.placeholder_with_default( 1., shape=()) # temperature for softmax self.edges_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes, vertexes)) self.nodes_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes)) self.embeddings = tf.placeholder(dtype=tf.float32, shape=(None, embedding_dim)) self.rewardR = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.rewardF = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.adjacency_tensor = tf.one_hot(self.edges_labels, depth=edges, dtype=tf.float32) self.node_tensor = tf.one_hot(self.nodes_labels, depth=nodes, dtype=tf.float32) with tf.variable_scope('generator'): self.edges_logits, self.nodes_logits = \ self.decoder(self.embeddings, decoder_units, vertexes, edges, nodes, training=self.training, dropout_rate=self.dropout_rate) with tf.name_scope('outputs'): (self.edges_softmax, self.nodes_softmax), \ (self.edges_argmax, self.nodes_argmax), \ (self.edges_gumbel_logits, self.nodes_gumbel_logits), \ (self.edges_gumbel_softmax, self.nodes_gumbel_softmax), \ (self.edges_gumbel_argmax, self.nodes_gumbel_argmax) = \ postprocess_logits((self.edges_logits, self.nodes_logits), temperature=self.temperature) self.edges_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.edges_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.edges_gumbel_argmax - self. edges_gumbel_softmax) + self. edges_gumbel_softmax }, default=lambda: self.edges_softmax, exclusive=True) self.nodes_hat = tf.case( { self.soft_gumbel_softmax: lambda: self.nodes_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient(self.nodes_gumbel_argmax - self. nodes_gumbel_softmax) + self. nodes_gumbel_softmax }, default=lambda: self.nodes_softmax, exclusive=True) with tf.name_scope('D_x_real'): self.logits_real, self.features_real = self.D_x( (self.adjacency_tensor, None, self.node_tensor), units=discriminator_units) with tf.name_scope('D_x_fake'): self.logits_fake, self.features_fake = self.D_x( (self.edges_hat, None, self.nodes_hat), units=discriminator_units) with tf.name_scope('V_x_real'): self.value_logits_real = self.V_x( (self.adjacency_tensor, None, self.node_tensor), units=discriminator_units) with tf.name_scope('V_x_fake'): self.value_logits_fake = self.V_x( (self.edges_hat, None, self.nodes_hat), units=discriminator_units)
def __init__(self, vertexes, edges, nodes, embedding_dim, decoder_units, discriminator_units, decoder, discriminator, soft_gumbel_softmax=False, hard_gumbel_softmax=False, batch_discriminator=True, unrolling_steps=1, batch_dim=32, latent_opt=False, noise_sigma=0.0, test_epoch=0): self.vertexes, self.edges, self.nodes, self.embedding_dim, self.decoder_units, self.discriminator_units, \ self.decoder, self.discriminator, self.batch_discriminator, self.unrolling_steps, self.latent_opt, self.noise_sigma, self.test_epoch = vertexes, edges, nodes, embedding_dim, decoder_units, discriminator_units, decoder, discriminator, batch_discriminator, unrolling_steps, latent_opt, noise_sigma, test_epoch self.training = tf.placeholder_with_default(False, shape=()) self.dropout_rate = tf.placeholder_with_default(0., shape=()) self.soft_gumbel_softmax = tf.placeholder_with_default(soft_gumbel_softmax, shape=()) self.hard_gumbel_softmax = tf.placeholder_with_default(hard_gumbel_softmax, shape=()) self.temperature = tf.placeholder_with_default(1., shape=()) self.edges_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes, vertexes)) self.nodes_labels = tf.placeholder(dtype=tf.int64, shape=(None, vertexes)) self.embeddings = tf.placeholder(dtype=tf.float32, shape=(None, embedding_dim)) self.input_len = tf.placeholder(dtype=tf.int64, shape=[]) self.is_training = True self.rewardR = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.rewardF = tf.placeholder(dtype=tf.float32, shape=(None, 1)) self.adjacency_tensor = tf.one_hot(self.edges_labels, depth=edges, dtype=tf.float32) self.node_tensor = tf.one_hot(self.nodes_labels, depth=nodes, dtype=tf.float32) print(f"self.is_training: {self.is_training}") with tf.variable_scope("input_LO"): # self.embeddings_LO = tf.Variable(tf.zeros(shape=(tf.shape(self.embeddings)[0], embedding_dim), dtype=tf.float32)) if self.is_training: self.embeddings_LO = tf.Variable(tf.zeros(shape=[batch_dim, embedding_dim]), name="latent_z") else: self.embeddings_LO = self.embeddings with tf.variable_scope('generator'): self.edges_logits, self.nodes_logits = self.decoder(self.embeddings_LO, decoder_units, vertexes, edges, nodes, training=self.training, dropout_rate=self.dropout_rate) with tf.name_scope('outputs'): (self.edges_softmax, self.nodes_softmax), \ (self.edges_argmax, self.nodes_argmax), \ (self.edges_gumbel_logits, self.nodes_gumbel_logits), \ (self.edges_gumbel_softmax, self.nodes_gumbel_softmax), \ (self.edges_gumbel_argmax, self.nodes_gumbel_argmax) = postprocess_logits( (self.edges_logits, self.nodes_logits), temperature=self.temperature) self.edges_hat = tf.case({self.soft_gumbel_softmax: lambda: self.edges_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient( self.edges_gumbel_argmax - self.edges_gumbel_softmax) + self.edges_gumbel_softmax}, default=lambda: self.edges_softmax, exclusive=True) self.nodes_hat = tf.case({self.soft_gumbel_softmax: lambda: self.nodes_gumbel_softmax, self.hard_gumbel_softmax: lambda: tf.stop_gradient( self.nodes_gumbel_argmax - self.nodes_gumbel_softmax) + self.nodes_gumbel_softmax}, default=lambda: self.nodes_softmax, exclusive=True) def Adding_Gaussian_Noise(input, noise_sigma=0.0): noise = tf.random_normal(shape=tf.shape(input), mean=0.0, stddev=noise_sigma, dtype=tf.float32) return input + noise with tf.name_scope('D_x_real'): self.logits_real, self.features_real = self.D_x((Adding_Gaussian_Noise(self.adjacency_tensor, noise_sigma=self.noise_sigma), None, Adding_Gaussian_Noise(self.node_tensor, noise_sigma=self.noise_sigma)), units=discriminator_units) with tf.name_scope('D_x_fake'): self.logits_fake, self.features_fake = self.D_x((Adding_Gaussian_Noise(self.edges_hat, noise_sigma=self.noise_sigma), None, Adding_Gaussian_Noise(self.nodes_hat, noise_sigma=self.noise_sigma)), units=discriminator_units) with tf.name_scope('V_x_real'): self.value_logits_real = self.V_x((self.adjacency_tensor, None, self.node_tensor), units=discriminator_units) with tf.name_scope('V_x_fake'): self.value_logits_fake = self.V_x((self.edges_hat, None, self.nodes_hat), units=discriminator_units)