Model documentation

demos/demo_gwhist_ggnn.py

@@ -129,7 +129,7 @@ def main():
                                               num_workers=args.prefetch, pin_memory=True)
 
     print('\tCreate model')
-    model = MpnnGGNN(stat_dict['edge_labels'], [len(h_t[0]), len(list(e.values())[0])], 25, 15, 2, num_classes, type='classification')
+    model = MpnnGGNN(stat_dict['edge_labels'], 25, 15, 2, num_classes, type='classification')
 
     print('Optimizer')
     optimizer = optim.Adam(model.parameters(), lr=args.lr)

demos/demo_letter_ggnn.py

@@ -130,7 +130,7 @@ def main():
                                               num_workers=args.prefetch, pin_memory=True)
 
     print('\tCreate model')
-    model = MpnnGGNN(stat_dict['edge_labels'], [len(h_t[0]), len(list(e.values())[0])], 25, 15, 2, num_classes, type='classification')
+    model = MpnnGGNN(stat_dict['edge_labels'], 25, 15, 2, num_classes, type='classification')
 
     print('Optimizer')
     optimizer = optim.Adam(model.parameters(), lr=args.lr)

demos/demo_qm9_ggnn.py

@@ -141,7 +141,7 @@ def main():
                                               num_workers=args.prefetch, pin_memory=True)
 
     print('\tCreate model')
-    model = MpnnGGNN(stat_dict['edge_labels'], [len(h_t[0]), len(list(e.values())[0])], 25, 15, 2, len(l),
+    model = MpnnGGNN(stat_dict['edge_labels'], 25, 15, 2, len(l),
                      type='regression')
 
     print('Optimizer')

models/MPNN.py

@@ -15,7 +15,22 @@
 
 class MPNN(nn.Module):
     """
-    in_n (size_v, size_e)
+        MPNN as proposed by Gilmer et al..
+
+        This class implements the whole Gilmer et al. model following the functions Message, Update and Readout.
+
+        Parameters
+        ----------
+        hidden_state_size : int
+            Size of the hidden states (the input will be padded with 0's to this size).
+        message_size : int
+            Message function output vector size.
+        n_layers : int
+            Number of iterations Message+Update (weight tying).
+        l_target : int
+            Size of the output.
+        type : str (Optional)
+            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
     """
 
     def __init__(self, in_n, hidden_state_size, message_size, n_layers, l_target, type='regression'):

models/MPNN_Duvenaud.py

@@ -15,8 +15,27 @@
 
 class MpnnDuvenaud(nn.Module):
     """
-    in_n (size_v, size_e)
+        MPNN as proposed by Duvenaud et al..
+
+        This class implements the whole Duvenaud et al. model following the functions proposed by Gilmer et al. as 
+        Message, Update and Readout.
+
+        Parameters
+        ----------
+        d : int list.
+            Possible degrees for the input graph.
+        in_n : int list
+            Sizes for the node and edge features.
+        out_update : int list
+            Output sizes for the different Update functions.
+        hidden_state_readout : int
+            Input size for the neural net used inside the readout function.
+        l_target : int
+            Size of the output.
+        type : str (Optional)
+            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
     """
+
     def __init__(self, d, in_n, out_update, hidden_state_readout, l_target, type='regression'):
         super(MpnnDuvenaud, self).__init__()
 

models/MPNN_GGNN.py

@@ -15,14 +15,32 @@
 
 class MpnnGGNN(nn.Module):
     """
-        in_n (size_v, size_e)
+        MPNN as proposed by Li et al..
+
+        This class implements the whole Li et al. model following the functions proposed by Gilmer et al. as
+        Message, Update and Readout.
+
+        Parameters
+        ----------
+        e : int list.
+            Possible edge labels for the input graph.
+        hidden_state_size : int
+            Size of the hidden states (the input will be padded with 0's to this size).
+        message_size : int
+            Message function output vector size.
+        n_layers : int
+            Number of iterations Message+Update (weight tying).
+        l_target : int
+            Size of the output.
+        type : str (Optional)
+            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
     """
 
-    def __init__(self, d, in_n, hidden_state_size, message_size, n_layers, l_target, type='regression'):
+    def __init__(self, e, hidden_state_size, message_size, n_layers, l_target, type='regression'):
         super(MpnnGGNN, self).__init__()
 
         # Define message
-        self.m = nn.ModuleList([MessageFunction('ggnn', args={'e_label': d, 'in': hidden_state_size, 'out': message_size})])
+        self.m = nn.ModuleList([MessageFunction('ggnn', args={'e_label': e, 'in': hidden_state_size, 'out': message_size})])
 
         # Define Update
         self.u = nn.ModuleList([UpdateFunction('ggnn',

models/MPNN_IntNet.py

@@ -15,8 +15,25 @@
 
 class MpnnIntNet(nn.Module):
     """
-    in_n (size_v, size_e)
+        MPNN as proposed by Battaglia et al..
+
+        This class implements the whole Battaglia et al. model following the functions proposed by Gilmer et al. as
+        Message, Update and Readout.
+
+        Parameters
+        ----------
+        in_n : int list
+            Sizes for the node and edge features.
+        out_message : int list
+            Output sizes for the different Message functions.
+        out_update : int list
+            Output sizes for the different Update functions.
+        l_target : int
+            Size of the output.
+        type : str (Optional)
+            Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
     """
+
     def __init__(self, in_n, out_message, out_update, l_target, type='regression'):
         super(MpnnIntNet, self).__init__()
 

models/README.md

@@ -1,24 +1,114 @@
 #Available MPNN models
 
-Explanation
+Some of the models available in the literature have been implemented as Message, Update and Readout functions.
 
-## Duvenaud
+## MpnnDuvenaud
 
-* Input:
-* Output:
+This class implements the whole Duvenaud et al. model following the functions proposed by Gilmer et al. as Message, Update and Readout.
 
-## GGNN
+``` 
+Parameters
+----------
+    d : int list.
+        Possible degrees for the input graph.
+    in_n : int list
+        Sizes for the node and edge features.
+    out_update : int list
+        Output sizes for the different Update funtion.
+    hidden_state_readout : int
+        Input size for the neural net used inside the readout function.
+    l_target : int
+        Size of the output.
+    type : str (Optional)
+        Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
+```
 
-* Input:
-* Output:
+Definition:
 
+``` 
+    model = MpnnDuvenaud(d, in_n, out_update, hidden_state_readout, l_target')
+```
+
+## MpnnGGNN
+
+
+This class implements the whole Li et al. model following the functions proposed by Gilmer et al. as Message, Update and Readout.
+
+``` 
+Parameters
+----------
+    e : int list.
+        Possible edge labels for the input graph.
+    hidden_state_size : int
+        Size of the hidden states (the input will be padded with 0's to this size).
+    message_size : int
+        Message function output vector size.
+    n_layers : int
+        Number of iterations Message+Update (weight tying).
+    l_target : int
+        Size of the output.
+    type : str (Optional)
+        Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
+```
+
+Definition:
+
+``` 
+    model = MpnnGGNN(e, in_n, hidden_state_size, message_size, n_layers, l_target)
+```
+
 ## IntNet
 
-* Input:
-* Output:
+This class implements the whole Battaglia et al. model following the functions proposed by Gilmer et al. as Message, Update and Readout.
+
+``` 
+Parameters
+----------
+    in_n : int list
+        Sizes for the node and edge features.
+    out_message : int list
+        Output sizes for the different Message functions.
+    out_update : int list
+        Output sizes for the different Update functions.
+    l_target : int
+        Size of the output.
+    type : str (Optional)
+        Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
+```
+
+Definition:
+
+``` 
+    model = MpnnIntNet(in_n, out_message, out_update, l_target):
+```
 
 ## MPNN as proposed by Gilmer et al.
 
-* Input:
-* Ouput:
+This class implements the whole Gilmer et al. model following the functions Message, Update and Readout.
+
+In progress..
+
+*[x] Edge Network
+*[ ] Virtual Graph Elements
+*[ ] set2set Readout function
+*[ ] Multiple Towers
+
+``` 
+Parameters
+----------
+    hidden_state_size : int
+        Size of the hidden states (the input will be padded with 0's to this size).
+    message_size : int
+        Message function output vector size.
+    n_layers : int
+        Number of iterations Message+Update (weight tying).
+    l_target : int
+        Size of the output.
+    type : str (Optional)
+        Classification | [Regression (default)]. If classification, LogSoftmax layer is applied to the output vector.
+``` 
+Definition:
 
+``` 
+    model = MPNN(in_n, hidden_state_size, message_size, n_layers, l_target, type='regression'):
+```