Updates to the MGN synthetic data reader

Author: szaman19

applications/graph/MeshGraphNet/GNN.py

@@ -1,5 +1,5 @@
 import lbann
-from .GNNComponents import MLP, GraphProcessor
+from GNNComponents import MLP, GraphProcessor
 
 
 def input_data_splitter(input_layer,
@@ -96,7 +96,7 @@ def LBANN_GNN_Model(num_nodes, num_edges,
 
   calculated_features = node_decoder(node_features)
 
-  loss = loss.MeanSquaredError(calculated_features, out_features)
+  loss = lbann.MeanSquaredError(calculated_features, out_features)
 
   # Define some of the usual callbacks
 

applications/graph/MeshGraphNet/GNNComponents.py

@@ -86,6 +86,7 @@ def __init__(self,
                norm_type=lbann.LayerNorm,
                name=None):
     super().__init__()
+    EdgeProcessor.global_count += 1
     self.instance = 0
     self.name = (name if name
                  else f'EdgeProcessor_{EdgeProcessor.global_count}')
@@ -139,6 +140,7 @@ def __init__(self,
                norm_type=lbann.LayerNorm,
                name=None):
     super().__init__()
+    NodeProcessor.global_count += 1
     self.instance = 0
     self.name = (name if name
                  else f'NodeProcessor_{NodeProcessor.global_count}')
@@ -166,7 +168,7 @@ def forward(self,
     self.instance += 1
 
     edge_feature_sum = lbann.Scatter(edge_features, target_edge_indices,
-                                     name="f{self.name}_{self.instance}_scatter",
+                                     name=f"{self.name}_{self.instance}_scatter",
                                      dims=[self.num_nodes, self.in_dim_edge],
                                      axis=0)
 
@@ -175,7 +177,7 @@ def forward(self,
                              name=f"{self.name}_{self.instance}_concat_features")
     x = self.node_mlp(x)
 
-    return lbann.Sum(edge_features, x,
+    return lbann.Sum(node_features, x,
                      name=f"{self.name}_{self.instance}_residual_sum")
 
 

applications/graph/MeshGraphNet/SyntheticData.py

@@ -1,12 +1,17 @@
 import numpy as np 
+import configparser
+
+
+DATA_CONFIG = configparser.ConfigParser()
+DATA_CONFIG.read("data_config.ini")
+NUM_NODES = 100 # int(DATA_CONFIG['DEFAULT']['NUM_NODES'])
+NUM_EDGES = 10000 # int(DATA_CONFIG['DEFAULT']['NUM_EDGES'])
+NODE_FEATS = 5 # int(DATA_CONFIG['DEFAULT']['NODE_FEATURES'])
+EDGE_FEATS = 3 # int(DATA_CONFIG['DEFAULT']['EDGE_FEATURES'])
+OUT_FEATS = 3 # int(DATA_CONFIG['DEFAULT']['OUT_FEATURES'])
+NUM_SAMPLES = 100
 
 
-NUM_SAMPLES = 10000
-NUM_NODES = 100
-NUM_EDGES = 1000
-NODE_FEATS = 5
-EDGE_FEATS = 3
-OUT_FEATS = 3
 
 NODE_FEATURE_SIZE = NUM_NODES * NODE_FEATS
 EDGE_FEATURE_SIZE = NUM_EDGES * EDGE_FEATS
@@ -27,3 +32,7 @@ def sample_dims_func():
 
   size = NODE_FEATURE_SIZE + EDGE_FEATURE_SIZE + OUT_FEATURE_SIZE + 2 * NUM_EDGES
   return (size, )
+
+
+if __name__ == '__main__':
+    print(NUM_NODES)

applications/graph/MeshGraphNet/Trainer.py

@@ -2,9 +2,9 @@
 import lbann.contrib.launcher
 import lbann.contrib.args
 import argparse
-
+import configparser
 import os.path as osp
-from .GNN import LBANN_GNN_Model
+from GNN import LBANN_GNN_Model
 
 data_dir = osp.dirname(osp.realpath(__file__))
 
@@ -21,8 +21,8 @@
     help='number of epochs (deafult: 3)', metavar='NUM')
 
 parser.add_argument(
-    '--mini-batch-size', action='store', default=256, type=int,
-    help="mini-batch size (default: 256)", metavar='NUM')
+    '--mini-batch-size', action='store', default=4, type=int,
+    help="mini-batch size (default: 4)", metavar='NUM')
 
 parser.add_argument(
     '--job-name', action='store', default="MGN", type=str,
@@ -38,12 +38,14 @@
 JOB_NAME = args.job_name
 
 # Some synthetic attributes to get the model running
+DATA_CONFIG = configparser.ConfigParser()
+DATA_CONFIG.read("data_config.ini")
 
-NUM_NODES = 100
-NUM_EDGES = 1000
-NODE_FEATS = 5
-EDGE_FEATS = 3
-OUT_FEATS = 3
+NUM_NODES = int(DATA_CONFIG['DEFAULT']['NUM_NODES'])
+NUM_EDGES = int(DATA_CONFIG['DEFAULT']['NUM_EDGES'])
+NODE_FEATS = int(DATA_CONFIG['DEFAULT']['NODE_FEATURES'])
+EDGE_FEATS = int(DATA_CONFIG['DEFAULT']['EDGE_FEATURES'])
+OUT_FEATS = int(DATA_CONFIG['DEFAULT']['OUT_FEATURES'])
 
 def make_data_reader(classname,
                      sample='get_sample_func',
@@ -85,4 +87,4 @@ def main():
                             **kwargs)
 
 if __name__ == '__main__':
-  main()
+  main()