Loading deepchem/models/layers.py +523 −18 Original line number Original line Diff line number Diff line Loading @@ -5,9 +5,9 @@ try: from collections.abc import Sequence as SequenceCollection from collections.abc import Sequence as SequenceCollection except: except: from collections import Sequence as SequenceCollection from collections import Sequence as SequenceCollection from typing import Callable, Dict, List from typing import Callable, Dict, List, Tuple from tensorflow.keras import activations, initializers, backend from tensorflow.keras import activations, initializers, backend from tensorflow.keras.layers import Dropout, BatchNormalization from tensorflow.keras.layers import Dropout, BatchNormalization, Dense, Activation class InteratomicL2Distances(tf.keras.layers.Layer): class InteratomicL2Distances(tf.keras.layers.Layer): Loading Loading @@ -370,6 +370,511 @@ class GraphGather(tf.keras.layers.Layer): return mol_features return mol_features class MolGANConvolutionLayer(tf.keras.layers.Layer): """ Graph convolution layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. Not used directly, higher level layers like MolGANMultiConvolutionLayer use it. This layer performs basic convolution on one-hot encoded matrices containing atom and bond information. This layer also accepts three inputs for the case when convolution is performed more than once and results of previous convolution need to used. It was done in such a way to avoid creating another layer that accepts three inputs rather than two. The last input layer is so-called hidden_layer and it hold results of the convolution while first two are unchanged input tensors. Example -------- See: MolGANMultiConvolutionLayer for using in layers. >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer1 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer2 = MolGANConvolutionLayer(units=units,edges=edges) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden1 = layer1([adjacency_tensor,node_tensor]) >>> output = layer2(hidden1) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: int, activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize this layer. Parameters --------- units: int Dimesion of dense layers used for convolution activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Dropout rate used by dropout layer edges: int, optional (default=5) How many dense layers to use in convolution. Typically equal to number of bond types used in the model. name: string, optional (default="") Name of the layer """ super(MolGANConvolutionLayer, self).__init__(name=name, **kwargs) self.activation = activation self.dropout_rate = dropout_rate self.units = units self.edges = edges self.dense1 = [Dense(units=self.units) for _ in range(edges - 1)] self.dense2 = Dense(units=self.units) self.dropout = Dropout(self.dropout_rate) self.activation_layer = Activation(self.activation) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- tuple(tf.Tensor,tf.Tensor,tf.Tensor) First and second are original input tensors Third is the result of convolution """ ic = len(inputs) assert ic > 1, "MolGANConvolutionLayer requires at least two inputs: [adjacency_tensor, node_features_tensor]" adjacency_tensor = inputs[0] node_tensor = inputs[1] # means that this is second loop of convolution if ic > 2: hidden_tensor = inputs[2] annotations = tf.concat((hidden_tensor, node_tensor), -1) else: annotations = node_tensor output = tf.stack([dense(annotations) for dense in self.dense1], 1) adj = tf.transpose(adjacency_tensor[:, :, :, 1:], (0, 3, 1, 2)) output = tf.matmul(adj, output) output = tf.reduce_sum(output, 1) + self.dense2(node_tensor) output = self.activation_layer(output) output = self.dropout(output) return adjacency_tensor, node_tensor, output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANConvolutionLayer, self).get_config() config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["units"] = self.units config["edges"] = self.edges return config class MolGANAggregationLayer(tf.keras.layers.Layer): """ Graph Aggregation layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. Performs aggregation on tensor resulting from convolution layers. Given its simple nature it might be removed in future and moved to MolGANEncoderLayer. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer_1 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer_2 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer_3 = MolGANAggregationLayer(units=128) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden_1 = layer_1([adjacency_tensor,node_tensor]) >>> hidden_2 = layer_2(hidden_1) >>> output = layer_3(hidden_2[2]) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) Example -------- vertices = 9 nodes = 5 edges = 5 units = 128 layer_1 = MolGANConvolutionLayer(units=units,edges=edges) layer_2 = MolGANConvolutionLayer(units=units,edges=edges) layer_3 = MolGANAggregationLayer(units=128) adjacency_tensor= layers.Input(shape=(vertices, vertices, edges)) node_tensor = layers.Input(shape=(vertices,nodes)) hidden_1 = layer_1([adjacency_tensor,node_tensor]) hidden_2 = layer_2(hidden_1) output = layer_3(hidden_2[2]) model = keras.Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: int = 128, activation: Callable = activations.tanh, dropout_rate: float = 0.0, name: str = "", **kwargs): """ Initialize the layer Parameters --------- units: int, optional (default=128) Dimesion of dense layers used for aggregation activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer name: string, optional (default="") Name of the layer """ super(MolGANAggregationLayer, self).__init__(name=name, **kwargs) self.units = units self.activation = activation self.dropout_rate = dropout_rate self.d1 = Dense(units=units, activation="sigmoid") self.d2 = Dense(units=units, activation=activation) self.dropout_layer = Dropout(dropout_rate) self.activation_layer = Activation(activation) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: List Single tensor resulting from graph convolution layer training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- aggregation tensor: tf.Tensor Result of aggregation function on input convolution tensor. """ i = self.d1(inputs) j = self.d2(inputs) output = tf.reduce_sum(i * j, 1) output = self.activation_layer(output) output = self.dropout_layer(output) return output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANAggregationLayer, self).get_config() config["units"] = self.units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class MolGANMultiConvolutionLayer(tf.keras.layers.Layer): """ Multiple pass convolution layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. It takes outputs of previous convolution layer and uses them as inputs for the next one. It simplifies the overall framework, but might be moved to MolGANEncoderLayer in the future in order to reduce number of layers. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer_1 = MolGANMultiConvolutionLayer(units=(128,64)) >>> layer_2 = MolGANAggregationLayer(units=128) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden = layer_1([adjacency_tensor,node_tensor]) >>> output = layer_2(hidden) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) Example -------- vertices = 9 nodes = 5 edges = 5 units = 128 layer_1 = MolGANMultiConvolutionLayer(units=(128,64)) layer_2 = MolGANAggregationLayer(units=128) adjacency_tensor= layers.Input(shape=(vertices, vertices, edges)) node_tensor = layers.Input(shape=(vertices,nodes)) hidden = layer_1([adjacency_tensor,node_tensor]) output = layer_2(hidden) model = keras.Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: Tuple = (128, 64), activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize the layer Parameters --------- units: Tuple, optional (default=(128,64)), min_length=2 List of dimensions used by consecutive convolution layers. The more values the more convolution layers invoked. activation: function, optional (default=tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer edges: int, optional (default=0) Controls how many dense layers use for single convolution unit. Typically matches number of bond types used in the molecule. name: string, optional (default="") Name of the layer """ super(MolGANMultiConvolutionLayer, self).__init__(name=name, **kwargs) assert len(units) > 1, "Layer requires at least two values" self.units = units self.activation = activation self.dropout_rate = dropout_rate self.edges = edges self.first_convolution = MolGANConvolutionLayer( self.units[0], self.activation, self.dropout_rate, self.edges) self.gcl = [ MolGANConvolutionLayer(u, self.activation, self.dropout_rate, self.edges) for u in self.units[1:] ] def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- convolution tensor: tf.Tensor Result of input tensors going through convolution a number of times. """ adjacency_tensor = inputs[0] node_tensor = inputs[1] tensors = self.first_convolution([adjacency_tensor, node_tensor]) for layer in self.gcl: tensors = layer(tensors) _, _, hidden_tensor = tensors return hidden_tensor def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANMultiConvolutionLayer, self).get_config() config["units"] = self.units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class MolGANEncoderLayer(tf.keras.layers.Layer): """ Main learning layer used by MolGAN model. MolGAN is a WGAN type model for generation of small molecules. It role is to further simplify model. This layer can be manually built by stacking graph convolution layers followed by graph aggregation. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input, Dropout,Dense >>> vertices = 9 >>> edges = 5 >>> nodes = 5 >>> dropout_rate = .0 >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices, nodes)) >>> graph = MolGANEncoderLayer(units = [(128,64),128], dropout_rate= dropout_rate, edges=edges)([adjacency_tensor,node_tensor]) >>> dense = Dense(units=128, activation='tanh')(graph) >>> dense = Dropout(dropout_rate)(dense) >>> dense = Dense(units=64, activation='tanh')(dense) >>> dense = Dropout(dropout_rate)(dense) >>> output = Dense(units=1)(dense) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: List = [(128, 64), 128], activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize the layer. Parameters --------- units: List, optional (default=[(128, 64), 128]) List of units for MolGANMultiConvolutionLayer and GraphAggregationLayer i.e. [(128,64),128] means two convolution layers dims = [128,64] followed by aggregation layer dims=128 activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer edges: int, optional (default=0) Controls how many dense layers use for single convolution unit. Typically matches number of bond types used in the molecule. name: string, optional (default="") Name of the layer """ super(MolGANEncoderLayer, self).__init__(name=name, **kwargs) assert len(units) == 2 self.graph_convolution_units, self.auxiliary_units = units self.activation = activation self.dropout_rate = dropout_rate self.edges = edges self.multi_graph_convolution_layer = MolGANMultiConvolutionLayer( self.graph_convolution_units, self.activation, self.dropout_rate, self.edges) self.graph_aggregation_layer = MolGANAggregationLayer( self.auxiliary_units, self.activation, self.dropout_rate) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- encoder tensor: tf.Tensor Tensor that been through number of convolutions followed by aggregation. """ output = self.multi_graph_convolution_layer(inputs) node_tensor = inputs[1] if len(inputs) > 2: hidden_tensor = inputs[2] annotations = tf.concat((output, hidden_tensor, node_tensor), -1) else: _, node_tensor = inputs annotations = tf.concat((output, node_tensor), -1) output = self.graph_aggregation_layer(annotations) return output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANEncoderLayer, self).get_config() config["graph_convolution_units"] = self.graph_convolution_units config["auxiliary_units"] = self.auxiliary_units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class LSTMStep(tf.keras.layers.Layer): class LSTMStep(tf.keras.layers.Layer): """Layer that performs a single step LSTM update. """Layer that performs a single step LSTM update. Loading deepchem/models/tests/test_molgan_layers.py 0 → 100644 +97 −0 Original line number Original line Diff line number Diff line import unittest from tensorflow import keras from tensorflow.keras.layers import Input from tensorflow.keras import activations from deepchem.models.layers import MolGANConvolutionLayer, MolGANMultiConvolutionLayer, MolGANAggregationLayer, MolGANEncoderLayer class test_molgan_layers(unittest.TestCase): """ Unit testing for MolGAN basic layers """ def test_graph_convolution_layer(self): vertices = 9 nodes = 5 edges = 5 units = 128 layer = MolGANConvolutionLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) output = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[output]) assert model.output_shape == [((None, vertices, vertices, edges), (None, vertices, nodes), (None, vertices, units))] assert layer.units == units assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 def test_aggregation_layer(self): vertices = 9 units = 128 layer = MolGANAggregationLayer(units=units) hidden_tensor = Input(shape=(vertices, units)) output = layer(hidden_tensor) model = keras.Model(inputs=[hidden_tensor], outputs=[output]) assert model.output_shape == (None, units) assert layer.units == units assert layer.activation == activations.tanh assert layer.dropout_rate == 0.0 def test_multigraph_convolution_layer(self): vertices = 9 nodes = 5 edges = 5 first_convolution_unit = 128 second_convolution_unit = 64 units = [first_convolution_unit, second_convolution_unit] layer = MolGANMultiConvolutionLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) hidden_tensor = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[hidden_tensor]) assert model.output_shape == (None, vertices, second_convolution_unit) assert layer.units == units assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 def test_graph_encoder_layer(self): vertices = 9 nodes = 5 edges = 5 first_convolution_unit = 128 second_convolution_unit = 64 aggregation_unit = 128 units = [(first_convolution_unit, second_convolution_unit), aggregation_unit] layer = MolGANEncoderLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) output = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[output]) assert model.output_shape == (None, aggregation_unit) assert layer.graph_convolution_units == (first_convolution_unit, second_convolution_unit) assert layer.auxiliary_units == aggregation_unit assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 if __name__ == '__main__': unittest.main() docs/source/api_reference/layers.rst +12 −0 Original line number Original line Diff line number Diff line Loading @@ -16,6 +16,18 @@ another tensor. DeepChem maintains an extensive collection of layers which perfo .. autoclass:: deepchem.models.layers.GraphGather .. autoclass:: deepchem.models.layers.GraphGather :members: :members: .. autoclass:: deepchem.models.layers.MolGANConvolutionLayer :members: .. autoclass:: deepchem.models.layers.MolGANAggregationLayer :members: .. autoclass:: deepchem.models.layers.MolGANMultiConvolutionLayer :members: .. autoclass:: deepchem.models.layers.MolGANEncoderLayer :members: .. autoclass:: deepchem.models.layers.LSTMStep .. autoclass:: deepchem.models.layers.LSTMStep :members: :members: Loading Loading
deepchem/models/layers.py +523 −18 Original line number Original line Diff line number Diff line Loading @@ -5,9 +5,9 @@ try: from collections.abc import Sequence as SequenceCollection from collections.abc import Sequence as SequenceCollection except: except: from collections import Sequence as SequenceCollection from collections import Sequence as SequenceCollection from typing import Callable, Dict, List from typing import Callable, Dict, List, Tuple from tensorflow.keras import activations, initializers, backend from tensorflow.keras import activations, initializers, backend from tensorflow.keras.layers import Dropout, BatchNormalization from tensorflow.keras.layers import Dropout, BatchNormalization, Dense, Activation class InteratomicL2Distances(tf.keras.layers.Layer): class InteratomicL2Distances(tf.keras.layers.Layer): Loading Loading @@ -370,6 +370,511 @@ class GraphGather(tf.keras.layers.Layer): return mol_features return mol_features class MolGANConvolutionLayer(tf.keras.layers.Layer): """ Graph convolution layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. Not used directly, higher level layers like MolGANMultiConvolutionLayer use it. This layer performs basic convolution on one-hot encoded matrices containing atom and bond information. This layer also accepts three inputs for the case when convolution is performed more than once and results of previous convolution need to used. It was done in such a way to avoid creating another layer that accepts three inputs rather than two. The last input layer is so-called hidden_layer and it hold results of the convolution while first two are unchanged input tensors. Example -------- See: MolGANMultiConvolutionLayer for using in layers. >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer1 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer2 = MolGANConvolutionLayer(units=units,edges=edges) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden1 = layer1([adjacency_tensor,node_tensor]) >>> output = layer2(hidden1) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: int, activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize this layer. Parameters --------- units: int Dimesion of dense layers used for convolution activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Dropout rate used by dropout layer edges: int, optional (default=5) How many dense layers to use in convolution. Typically equal to number of bond types used in the model. name: string, optional (default="") Name of the layer """ super(MolGANConvolutionLayer, self).__init__(name=name, **kwargs) self.activation = activation self.dropout_rate = dropout_rate self.units = units self.edges = edges self.dense1 = [Dense(units=self.units) for _ in range(edges - 1)] self.dense2 = Dense(units=self.units) self.dropout = Dropout(self.dropout_rate) self.activation_layer = Activation(self.activation) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- tuple(tf.Tensor,tf.Tensor,tf.Tensor) First and second are original input tensors Third is the result of convolution """ ic = len(inputs) assert ic > 1, "MolGANConvolutionLayer requires at least two inputs: [adjacency_tensor, node_features_tensor]" adjacency_tensor = inputs[0] node_tensor = inputs[1] # means that this is second loop of convolution if ic > 2: hidden_tensor = inputs[2] annotations = tf.concat((hidden_tensor, node_tensor), -1) else: annotations = node_tensor output = tf.stack([dense(annotations) for dense in self.dense1], 1) adj = tf.transpose(adjacency_tensor[:, :, :, 1:], (0, 3, 1, 2)) output = tf.matmul(adj, output) output = tf.reduce_sum(output, 1) + self.dense2(node_tensor) output = self.activation_layer(output) output = self.dropout(output) return adjacency_tensor, node_tensor, output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANConvolutionLayer, self).get_config() config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["units"] = self.units config["edges"] = self.edges return config class MolGANAggregationLayer(tf.keras.layers.Layer): """ Graph Aggregation layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. Performs aggregation on tensor resulting from convolution layers. Given its simple nature it might be removed in future and moved to MolGANEncoderLayer. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer_1 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer_2 = MolGANConvolutionLayer(units=units,edges=edges) >>> layer_3 = MolGANAggregationLayer(units=128) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden_1 = layer_1([adjacency_tensor,node_tensor]) >>> hidden_2 = layer_2(hidden_1) >>> output = layer_3(hidden_2[2]) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) Example -------- vertices = 9 nodes = 5 edges = 5 units = 128 layer_1 = MolGANConvolutionLayer(units=units,edges=edges) layer_2 = MolGANConvolutionLayer(units=units,edges=edges) layer_3 = MolGANAggregationLayer(units=128) adjacency_tensor= layers.Input(shape=(vertices, vertices, edges)) node_tensor = layers.Input(shape=(vertices,nodes)) hidden_1 = layer_1([adjacency_tensor,node_tensor]) hidden_2 = layer_2(hidden_1) output = layer_3(hidden_2[2]) model = keras.Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: int = 128, activation: Callable = activations.tanh, dropout_rate: float = 0.0, name: str = "", **kwargs): """ Initialize the layer Parameters --------- units: int, optional (default=128) Dimesion of dense layers used for aggregation activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer name: string, optional (default="") Name of the layer """ super(MolGANAggregationLayer, self).__init__(name=name, **kwargs) self.units = units self.activation = activation self.dropout_rate = dropout_rate self.d1 = Dense(units=units, activation="sigmoid") self.d2 = Dense(units=units, activation=activation) self.dropout_layer = Dropout(dropout_rate) self.activation_layer = Activation(activation) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: List Single tensor resulting from graph convolution layer training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- aggregation tensor: tf.Tensor Result of aggregation function on input convolution tensor. """ i = self.d1(inputs) j = self.d2(inputs) output = tf.reduce_sum(i * j, 1) output = self.activation_layer(output) output = self.dropout_layer(output) return output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANAggregationLayer, self).get_config() config["units"] = self.units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class MolGANMultiConvolutionLayer(tf.keras.layers.Layer): """ Multiple pass convolution layer used in MolGAN model. MolGAN is a WGAN type model for generation of small molecules. It takes outputs of previous convolution layer and uses them as inputs for the next one. It simplifies the overall framework, but might be moved to MolGANEncoderLayer in the future in order to reduce number of layers. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input >>> vertices = 9 >>> nodes = 5 >>> edges = 5 >>> units = 128 >>> layer_1 = MolGANMultiConvolutionLayer(units=(128,64)) >>> layer_2 = MolGANAggregationLayer(units=128) >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices,nodes)) >>> hidden = layer_1([adjacency_tensor,node_tensor]) >>> output = layer_2(hidden) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) Example -------- vertices = 9 nodes = 5 edges = 5 units = 128 layer_1 = MolGANMultiConvolutionLayer(units=(128,64)) layer_2 = MolGANAggregationLayer(units=128) adjacency_tensor= layers.Input(shape=(vertices, vertices, edges)) node_tensor = layers.Input(shape=(vertices,nodes)) hidden = layer_1([adjacency_tensor,node_tensor]) output = layer_2(hidden) model = keras.Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: Tuple = (128, 64), activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize the layer Parameters --------- units: Tuple, optional (default=(128,64)), min_length=2 List of dimensions used by consecutive convolution layers. The more values the more convolution layers invoked. activation: function, optional (default=tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer edges: int, optional (default=0) Controls how many dense layers use for single convolution unit. Typically matches number of bond types used in the molecule. name: string, optional (default="") Name of the layer """ super(MolGANMultiConvolutionLayer, self).__init__(name=name, **kwargs) assert len(units) > 1, "Layer requires at least two values" self.units = units self.activation = activation self.dropout_rate = dropout_rate self.edges = edges self.first_convolution = MolGANConvolutionLayer( self.units[0], self.activation, self.dropout_rate, self.edges) self.gcl = [ MolGANConvolutionLayer(u, self.activation, self.dropout_rate, self.edges) for u in self.units[1:] ] def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- convolution tensor: tf.Tensor Result of input tensors going through convolution a number of times. """ adjacency_tensor = inputs[0] node_tensor = inputs[1] tensors = self.first_convolution([adjacency_tensor, node_tensor]) for layer in self.gcl: tensors = layer(tensors) _, _, hidden_tensor = tensors return hidden_tensor def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANMultiConvolutionLayer, self).get_config() config["units"] = self.units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class MolGANEncoderLayer(tf.keras.layers.Layer): """ Main learning layer used by MolGAN model. MolGAN is a WGAN type model for generation of small molecules. It role is to further simplify model. This layer can be manually built by stacking graph convolution layers followed by graph aggregation. Example -------- >>> from tensorflow.keras import Model >>> from tensorflow.keras.layers import Input, Dropout,Dense >>> vertices = 9 >>> edges = 5 >>> nodes = 5 >>> dropout_rate = .0 >>> adjacency_tensor= Input(shape=(vertices, vertices, edges)) >>> node_tensor = Input(shape=(vertices, nodes)) >>> graph = MolGANEncoderLayer(units = [(128,64),128], dropout_rate= dropout_rate, edges=edges)([adjacency_tensor,node_tensor]) >>> dense = Dense(units=128, activation='tanh')(graph) >>> dense = Dropout(dropout_rate)(dense) >>> dense = Dense(units=64, activation='tanh')(dense) >>> dense = Dropout(dropout_rate)(dense) >>> output = Dense(units=1)(dense) >>> model = Model(inputs=[adjacency_tensor,node_tensor], outputs=[output]) References ---------- .. [1] Nicola De Cao et al. "MolGAN: An implicit generative model for small molecular graphs", https://arxiv.org/abs/1805.11973 """ def __init__(self, units: List = [(128, 64), 128], activation: Callable = activations.tanh, dropout_rate: float = 0.0, edges: int = 5, name: str = "", **kwargs): """ Initialize the layer. Parameters --------- units: List, optional (default=[(128, 64), 128]) List of units for MolGANMultiConvolutionLayer and GraphAggregationLayer i.e. [(128,64),128] means two convolution layers dims = [128,64] followed by aggregation layer dims=128 activation: function, optional (default=Tanh) activation function used across model, default is Tanh dropout_rate: float, optional (default=0.0) Used by dropout layer edges: int, optional (default=0) Controls how many dense layers use for single convolution unit. Typically matches number of bond types used in the molecule. name: string, optional (default="") Name of the layer """ super(MolGANEncoderLayer, self).__init__(name=name, **kwargs) assert len(units) == 2 self.graph_convolution_units, self.auxiliary_units = units self.activation = activation self.dropout_rate = dropout_rate self.edges = edges self.multi_graph_convolution_layer = MolGANMultiConvolutionLayer( self.graph_convolution_units, self.activation, self.dropout_rate, self.edges) self.graph_aggregation_layer = MolGANAggregationLayer( self.auxiliary_units, self.activation, self.dropout_rate) def call(self, inputs, training=False): """ Invoke this layer Parameters ---------- inputs: list List of two input matrices, adjacency tensor and node features tensors in one-hot encoding format. training: bool Should this layer be run in training mode. Typically decided by main model, influences things like dropout. Returns -------- encoder tensor: tf.Tensor Tensor that been through number of convolutions followed by aggregation. """ output = self.multi_graph_convolution_layer(inputs) node_tensor = inputs[1] if len(inputs) > 2: hidden_tensor = inputs[2] annotations = tf.concat((output, hidden_tensor, node_tensor), -1) else: _, node_tensor = inputs annotations = tf.concat((output, node_tensor), -1) output = self.graph_aggregation_layer(annotations) return output def get_config(self) -> Dict: """ Returns config dictionary for this layer. """ config = super(MolGANEncoderLayer, self).get_config() config["graph_convolution_units"] = self.graph_convolution_units config["auxiliary_units"] = self.auxiliary_units config["activation"] = self.activation config["dropout_rate"] = self.dropout_rate config["edges"] = self.edges return config class LSTMStep(tf.keras.layers.Layer): class LSTMStep(tf.keras.layers.Layer): """Layer that performs a single step LSTM update. """Layer that performs a single step LSTM update. Loading
deepchem/models/tests/test_molgan_layers.py 0 → 100644 +97 −0 Original line number Original line Diff line number Diff line import unittest from tensorflow import keras from tensorflow.keras.layers import Input from tensorflow.keras import activations from deepchem.models.layers import MolGANConvolutionLayer, MolGANMultiConvolutionLayer, MolGANAggregationLayer, MolGANEncoderLayer class test_molgan_layers(unittest.TestCase): """ Unit testing for MolGAN basic layers """ def test_graph_convolution_layer(self): vertices = 9 nodes = 5 edges = 5 units = 128 layer = MolGANConvolutionLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) output = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[output]) assert model.output_shape == [((None, vertices, vertices, edges), (None, vertices, nodes), (None, vertices, units))] assert layer.units == units assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 def test_aggregation_layer(self): vertices = 9 units = 128 layer = MolGANAggregationLayer(units=units) hidden_tensor = Input(shape=(vertices, units)) output = layer(hidden_tensor) model = keras.Model(inputs=[hidden_tensor], outputs=[output]) assert model.output_shape == (None, units) assert layer.units == units assert layer.activation == activations.tanh assert layer.dropout_rate == 0.0 def test_multigraph_convolution_layer(self): vertices = 9 nodes = 5 edges = 5 first_convolution_unit = 128 second_convolution_unit = 64 units = [first_convolution_unit, second_convolution_unit] layer = MolGANMultiConvolutionLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) hidden_tensor = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[hidden_tensor]) assert model.output_shape == (None, vertices, second_convolution_unit) assert layer.units == units assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 def test_graph_encoder_layer(self): vertices = 9 nodes = 5 edges = 5 first_convolution_unit = 128 second_convolution_unit = 64 aggregation_unit = 128 units = [(first_convolution_unit, second_convolution_unit), aggregation_unit] layer = MolGANEncoderLayer(units=units, edges=edges) adjacency_tensor = Input(shape=(vertices, vertices, edges)) node_tensor = Input(shape=(vertices, nodes)) output = layer([adjacency_tensor, node_tensor]) model = keras.Model( inputs=[adjacency_tensor, node_tensor], outputs=[output]) assert model.output_shape == (None, aggregation_unit) assert layer.graph_convolution_units == (first_convolution_unit, second_convolution_unit) assert layer.auxiliary_units == aggregation_unit assert layer.activation == activations.tanh assert layer.edges == 5 assert layer.dropout_rate == 0.0 if __name__ == '__main__': unittest.main()
docs/source/api_reference/layers.rst +12 −0 Original line number Original line Diff line number Diff line Loading @@ -16,6 +16,18 @@ another tensor. DeepChem maintains an extensive collection of layers which perfo .. autoclass:: deepchem.models.layers.GraphGather .. autoclass:: deepchem.models.layers.GraphGather :members: :members: .. autoclass:: deepchem.models.layers.MolGANConvolutionLayer :members: .. autoclass:: deepchem.models.layers.MolGANAggregationLayer :members: .. autoclass:: deepchem.models.layers.MolGANMultiConvolutionLayer :members: .. autoclass:: deepchem.models.layers.MolGANEncoderLayer :members: .. autoclass:: deepchem.models.layers.LSTMStep .. autoclass:: deepchem.models.layers.LSTMStep :members: :members: Loading
