Merge pull request #386 from corochann/graph-film-v2

Graph film v2

Author: mottodora

README.md

@@ -111,6 +111,7 @@ The following graph convolutional neural networks are currently supported:
 - GIN: Graph Isomorphism Networks [17]
 - MPNN: Message Passing Neural Networks [3]
 - Set2Set [19]
+- GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation [20]
 
 We test supporting the brand-new Graph Warp Module (GWM) [18]-attached models for:
 - NFP ('nfp_gwm')
@@ -202,3 +203,6 @@ papers. Use the library at your own risk.
 [18] K. Ishiguro, S. Maeda, and M. Koyama, ``Graph Warp Module: an Auxiliary Module for Boosting the Power of Graph Neural Networks'', arXiv:1902.01020 [cs.LG], 2019.
 
 [19] Oriol Vinyals, Samy Bengio, Manjunath Kudlur. Order Matters: Sequence to sequence for sets. *arXiv preprint arXiv:1511.06391*, 2015.
+. 
+
+[20] Marc Brockschmidt, ``GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation'', arXiv:1906.12192 [cs.ML], 2019.

chainer_chemistry/__init__.py

@@ -3,13 +3,16 @@
 from chainer_chemistry import dataset  # NOQA
 try:
     from chainer_chemistry import datasets  # NOQA
-except ImportError:
-    warnings.warn(
-        'A module chainer_chemistry.datasets was not imported, '
-        'probably because RDKit is not installed. '
-        'To install RDKit, please follow instruction in '
-        'https://github.com/pfnet-research/chainer-chemistry#installation.',
-        UserWarning)
+except ImportError as e:
+    if 'rdkit' in e.msg:
+        warnings.warn(
+            'A module chainer_chemistry.datasets was not imported, '
+            'probably because RDKit is not installed. '
+            'To install RDKit, please follow instruction in '
+            'https://github.com/pfnet-research/chainer-chemistry#installation.',
+            UserWarning)
+    else:
+        raise(e)
 from chainer_chemistry import functions  # NOQA
 from chainer_chemistry import links  # NOQA
 from chainer_chemistry import models  # NOQA

chainer_chemistry/dataset/preprocessors/__init__.py

@@ -9,6 +9,7 @@
 from chainer_chemistry.dataset.preprocessors.ecfp_preprocessor import ECFPPreprocessor  # NOQA
 from chainer_chemistry.dataset.preprocessors.relgat_preprocessor import RelGATPreprocessor  # NOQA
 from chainer_chemistry.dataset.preprocessors.ggnn_preprocessor import GGNNPreprocessor  # NOQA
+from chainer_chemistry.dataset.preprocessors.gnnfilm_preprocessor import GNNFiLMPreprocessor  # NOQA
 from chainer_chemistry.dataset.preprocessors.gin_preprocessor import GINPreprocessor  # NOQA
 from chainer_chemistry.dataset.preprocessors.gwm_preprocessor import GGNNGWMPreprocessor  # NOQA
 from chainer_chemistry.dataset.preprocessors.gwm_preprocessor import GINGWMPreprocessor  # NOQA
@@ -35,4 +36,5 @@
     'rsgcn': RSGCNPreprocessor,
     'rsgcn_gwm': RSGCNGWMPreprocessor,
     'relgat': RelGATPreprocessor,
+    'gnnfilm': GNNFiLMPreprocessor,
 }

chainer_chemistry/dataset/preprocessors/common.py

@@ -118,7 +118,8 @@ def construct_adj_matrix(mol, out_size=-1, self_connection=True):
     return adj_array
 
 
-def construct_discrete_edge_matrix(mol, out_size=-1):
+def construct_discrete_edge_matrix(mol, out_size=-1,
+                                   add_self_connection_channel=False):
     """Returns the edge-type dependent adjacency matrix of the given molecule.
 
     Args:
@@ -129,6 +130,9 @@ def construct_discrete_edge_matrix(mol, out_size=-1):
             in the input molecules. In that case, the adjacent
             matrix is expanded and zeros are padded to right
             columns and bottom rows.
+        add_self_connection_channel (bool): Add self connection or not.
+            If True, adjacency matrix whose diagonal element filled with 1
+            is added to last channel.
 
     Returns:
         adj_array (numpy.ndarray): The adjacent matrix of the input molecule.
@@ -150,7 +154,10 @@ def construct_discrete_edge_matrix(mol, out_size=-1):
         raise ValueError(
             'out_size {} is smaller than number of atoms in mol {}'
             .format(out_size, N))
-    adjs = numpy.zeros((4, size, size), dtype=numpy.float32)
+    if add_self_connection_channel:
+        adjs = numpy.zeros((5, size, size), dtype=numpy.float32)
+    else:
+        adjs = numpy.zeros((4, size, size), dtype=numpy.float32)
 
     bond_type_to_channel = {
         Chem.BondType.SINGLE: 0,
@@ -165,6 +172,8 @@ def construct_discrete_edge_matrix(mol, out_size=-1):
         j = bond.GetEndAtomIdx()
         adjs[ch, i, j] = 1.0
         adjs[ch, j, i] = 1.0
+    if add_self_connection_channel:
+        adjs[-1] = numpy.eye(N)
     return adjs
 
 

chainer_chemistry/dataset/preprocessors/gin_preprocessor.py

@@ -32,7 +32,6 @@ def __init__(self, max_atoms=-1, out_size=-1, add_Hs=False):
         self.max_atoms = max_atoms
         self.out_size = out_size
 
-
     def get_input_features(self, mol):
         """get input features
 
@@ -45,4 +44,4 @@ def get_input_features(self, mol):
         type_check_num_atoms(mol, self.max_atoms)
         atom_array = construct_atomic_number_array(mol, out_size=self.out_size)
         adj_array = construct_adj_matrix(mol, out_size=self.out_size)
-        return atom_array, adj_array
+        return atom_array, adj_array

chainer_chemistry/dataset/preprocessors/gnnfilm_preprocessor.py

@@ -0,0 +1,48 @@
+from chainer_chemistry.dataset.preprocessors.common \
+    import construct_atomic_number_array, construct_discrete_edge_matrix
+from chainer_chemistry.dataset.preprocessors.common import type_check_num_atoms
+from chainer_chemistry.dataset.preprocessors.mol_preprocessor \
+    import MolPreprocessor
+
+
+class GNNFiLMPreprocessor(MolPreprocessor):
+    """GNNFiLM Preprocessor
+
+    Args:
+        max_atoms (int): Max number of atoms for each molecule, if the
+            number of atoms is more than this value, this data is simply
+            ignored.
+            Setting negative value indicates no limit for max atoms.
+        out_size (int): It specifies the size of array returned by
+            `get_input_features`.
+            If the number of atoms in the molecule is less than this value,
+            the returned arrays is padded to have fixed size.
+            Setting negative value indicates do not pad returned array.
+        add_Hs (bool): If True, implicit Hs are added.
+        kekulize (bool): If True, Kekulizes the molecule.
+
+    """
+
+    def __init__(self, max_atoms=-1, out_size=-1, add_Hs=False,
+                 kekulize=False):
+        super(GNNFiLMPreprocessor, self).__init__(
+            add_Hs=add_Hs, kekulize=kekulize)
+        if max_atoms >= 0 and out_size >= 0 and max_atoms > out_size:
+            raise ValueError('max_atoms {} must be less or equal to '
+                             'out_size {}'.format(max_atoms, out_size))
+        self.max_atoms = max_atoms
+        self.out_size = out_size
+
+    def get_input_features(self, mol):
+        """get input features
+
+        Args:
+            mol (Mol): Molecule input
+
+        Returns:
+
+        """
+        type_check_num_atoms(mol, self.max_atoms)
+        atom_array = construct_atomic_number_array(mol, out_size=self.out_size)
+        adj_array = construct_discrete_edge_matrix(mol, out_size=self.out_size, add_self_connection_channel=True)
+        return atom_array, adj_array

chainer_chemistry/dataset/preprocessors/mol_preprocessor.py

@@ -39,7 +39,6 @@ def prepare_smiles_and_mol(self, mol):
             Chem.Kekulize(mol)
         return canonical_smiles, mol
 
-
     def get_label(self, mol, label_names=None):
         """Extracts label information from a molecule.
 

chainer_chemistry/links/update/gnn_film_update.py

@@ -0,0 +1,59 @@
+import chainer
+from chainer import functions
+from chainer import links
+
+from chainer_chemistry.links.connection.graph_linear import GraphLinear
+
+
+class GNNFiLMUpdate(chainer.Chain):
+    """GNNFiLM submodule for update part.
+
+    Args:
+        hidden_channels (int): dimension of feature vector associated to
+            each atom
+        n_edge_types (int): number of types of edge
+    """
+
+    def __init__(self, hidden_channels=16, n_edge_types=5, activation=functions.relu):
+        super(GNNFiLMUpdate, self).__init__()
+        self.n_edge_types = n_edge_types
+        self.activation = activation
+        with self.init_scope():
+            self.W_linear = GraphLinear(
+                in_size=None, out_size=self.n_edge_types * hidden_channels, nobias=True)  # W_l in eq. (6)
+            self.W_g = GraphLinear(
+                in_size=None, out_size=self.n_edge_types * hidden_channels * 2, nobias=True)  # g in eq. (6)
+            self.norm_layer = links.LayerNormalization()  # l in eq. (6)
+
+    def forward(self, h, adj):
+        # --- Message part ---
+
+        xp = self.xp
+        mb, atom, ch = h.shape
+        newshape = adj.shape + (ch, )
+        adj = functions.broadcast_to(adj[:, :, :, :, xp.newaxis], newshape)
+        messages = functions.reshape(self.W_linear(h),
+                                     (mb, atom, ch, self.n_edge_types))
+        messages = functions.transpose(messages, (3, 0, 1, 2))
+        film_weights = functions.reshape(self.W_g(h),
+                                         (mb, atom, 2 * ch, self.n_edge_types))
+        film_weights = functions.transpose(film_weights, (3, 0, 1, 2))
+        # (n_edge_types, minibatch, atom, out_ch)
+        gamma = film_weights[:, :, :, :ch]
+        # (n_edge_types, minibatch, atom, out_ch)
+        beta = film_weights[:, :, :, ch:]
+
+        # --- Update part ---
+
+        messages = functions.expand_dims(gamma, axis=3) * functions.expand_dims(
+            messages, axis=2) + functions.expand_dims(beta, axis=3)
+        messages = self.activation(messages)
+        # (minibatch, n_edge_types, atom, atom, out_ch)
+        messages = functions.transpose(messages, (1, 0, 2, 3, 4))
+        messages = adj * messages
+        messages = functions.sum(messages, axis=3)  # sum across atoms
+        messages = functions.sum(messages, axis=1)  # sum across n_edge_types
+        messages = functions.reshape(messages, (mb * atom, ch))
+        messages = self.norm_layer(messages)
+        messages = functions.reshape(messages, (mb, atom, ch))
+        return messages

chainer_chemistry/models/__init__.py

@@ -22,6 +22,7 @@
 from chainer_chemistry.models.rsgcn import RSGCN  # NOQA
 from chainer_chemistry.models.schnet import SchNet  # NOQA
 from chainer_chemistry.models.weavenet import WeaveNet  # NOQA
+from chainer_chemistry.models.gnn_film import GNNFiLM  # NOQA
 
 from chainer_chemistry.models.gwm.gwm_net import GGNN_GWM  # NOQA
 from chainer_chemistry.models.gwm.gwm_net import GIN_GWM  # NOQA

chainer_chemistry/models/gnn_film.py

@@ -0,0 +1,104 @@
+import chainer
+from chainer import cuda
+from chainer import functions
+
+from chainer_chemistry.config import MAX_ATOMIC_NUM
+from chainer_chemistry.links.connection.embed_atom_id import EmbedAtomID
+from chainer_chemistry.links.readout.ggnn_readout import GGNNReadout
+from chainer_chemistry.links.update.gnn_film_update import GNNFiLMUpdate
+
+
+class GNNFiLM(chainer.Chain):
+    """Graph Neural Networks with Feature-wise Linear Modulation (GNN_FiLM)
+
+    Marc Brockschmidt (2019).\
+        GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation \
+        `arXiv:1906.12192 <https://arxiv.org/abs/1906.12192>`_
+
+    Args:
+        out_dim (int): dimension of output feature vector
+        hidden_channels (int): dimension of feature vector
+            associated to each atom
+        n_update_layers (int): number of layers
+        n_atom_types (int): number of types of atoms
+        concat_hidden (bool): If set to True, readout is executed in each layer
+            and the result is concatenated
+        weight_tying (bool): enable weight_tying or not
+        activation (~chainer.Function or ~chainer.FunctionNode):
+            activate function
+        n_edge_types (int): number of edge type.
+            Defaults to 5 for single, double, triple, aromatic bond
+            and self-connection.
+    """
+
+    def __init__(self, out_dim, hidden_channels=16, n_update_layers=4,
+                 n_atom_types=MAX_ATOMIC_NUM, concat_hidden=False,
+                 weight_tying=True, activation=functions.identity,
+                 n_edge_types=5):
+        super(GNNFiLM, self).__init__()
+        n_readout_layer = n_update_layers if concat_hidden else 1
+        n_message_layer = 1 if weight_tying else n_update_layers
+        with self.init_scope():
+            # Update
+            self.embed = EmbedAtomID(out_size=hidden_channels,
+                                     in_size=n_atom_types)
+            self.update_layers = chainer.ChainList(*[GNNFiLMUpdate(
+                hidden_channels=hidden_channels, n_edge_types=n_edge_types)
+                for _ in range(n_message_layer)])
+            # Readout
+            # self.readout_layers = chainer.ChainList(*[GeneralReadout(
+            #     out_dim=out_dim, hidden_channels=hidden_channels,
+            #     activation=activation, activation_agg=activation)
+            #     for _ in range(n_readout_layer)])
+            self.readout_layers = chainer.ChainList(*[GGNNReadout(
+                out_dim=out_dim, in_channels=hidden_channels * 2,
+                activation=activation, activation_agg=activation)
+                for _ in range(n_readout_layer)])
+        self.out_dim = out_dim
+        self.hidden_channels = hidden_channels
+        self.n_update_layers = n_update_layers
+        self.n_edge_types = n_edge_types
+        self.activation = activation
+        self.concat_hidden = concat_hidden
+        self.weight_tying = weight_tying
+
+    def __call__(self, atom_array, adj, is_real_node=None):
+        """Forward propagation
+
+        Args:
+            atom_array (numpy.ndarray): minibatch of molecular which is
+                represented with atom IDs (representing C, O, S, ...)
+                `atom_array[mol_index, atom_index]` represents `mol_index`-th
+                molecule's `atom_index`-th atomic number
+            adj (numpy.ndarray): minibatch of adjancency matrix with edge-type
+                information
+            is_real_node (numpy.ndarray): 2-dim array (minibatch, num_nodes).
+                1 for real node, 0 for virtual node.
+                If `None`, all node is considered as real node.
+
+        Returns:
+            ~chainer.Variable: minibatch of fingerprint
+        """
+        # reset state
+        # self.reset_state()
+        if atom_array.dtype == self.xp.int32:
+            h = self.embed(atom_array)  # (minibatch, max_num_atoms)
+        else:
+            h = atom_array
+        h0 = functions.copy(h, cuda.get_device_from_array(h.data).id)
+        g_list = []
+        for step in range(self.n_update_layers):
+            message_layer_index = 0 if self.weight_tying else step
+            h = self.update_layers[message_layer_index](h, adj)
+            if self.concat_hidden:
+                g = self.readout_layers[step](h, h0, is_real_node)
+                g_list.append(g)
+
+        if self.concat_hidden:
+            return functions.concat(g_list, axis=1)
+        else:
+            g = self.readout_layers[0](h, h0, is_real_node)
+            return g
+
+    def reset_state(self):
+        [update_layer.reset_state() for update_layer in self.update_layers]