Merge pull request #156 from ModECI/development

Development
ModECI · Oct 11, 2021 · 8facbb6 · 8facbb6
2 parents 77a6c6c + 8ecbabe
commit 8facbb6
Show file tree

Hide file tree

Showing 10 changed files with 555 additions and 17 deletions.
diff --git a/examples/ONNX/ab.png b/examples/ONNX/ab.png
diff --git a/examples/ONNX/ab_torch-jit-export-m2o.onnx b/examples/ONNX/ab_torch-jit-export-m2o.onnx
diff --git a/examples/ONNX/abc.png b/examples/ONNX/abc.png
diff --git a/examples/ONNX/abc_basic-mdf.json b/examples/ONNX/abc_basic-mdf.json
@@ -0,0 +1,86 @@
+{
+    "ONNX Model": {
+        "format": "ModECI MDF v0.2",
+        "generating_application": "Python modeci-mdf v0.2.1",
+        "graphs": {
+            "torch-jit-export": {
+                "nodes": {
+                    "Sin_0": {
+                        "input_ports": {
+                            "input": {
+                                "shape": "(2, 3)"
+                            }
+                        },
+                        "parameters": {
+                            "Sin_0": {
+                                "function": "onnx::Sin",
+                                "args": {
+                                    "input": "input"
+                                }
+                            }
+                        },
+                        "output_ports": {
+                            "_1": {
+                                "value": "Sin_0"
+                            }
+                        }
+                    },
+                    "Sin_1": {
+                        "input_ports": {
+                            "_1": {
+                                "shape": "(2, 3)"
+                            }
+                        },
+                        "parameters": {
+                            "Sin_1": {
+                                "function": "onnx::Sin",
+                                "args": {
+                                    "input": "_1"
+                                }
+                            }
+                        },
+                        "output_ports": {
+                            "_2": {
+                                "value": "Sin_1"
+                            }
+                        }
+                    },
+                    "Cos_2": {
+                        "input_ports": {
+                            "_2": {
+                                "shape": "(2, 3)"
+                            }
+                        },
+                        "parameters": {
+                            "Cos_2": {
+                                "function": "onnx::Cos",
+                                "args": {
+                                    "input": "_2"
+                                }
+                            }
+                        },
+                        "output_ports": {
+                            "_3": {
+                                "value": "Cos_2"
+                            }
+                        }
+                    }
+                },
+                "edges": {
+                    "Sin_0._1_Sin_1._1": {
+                        "sender": "Sin_0",
+                        "receiver": "Sin_1",
+                        "sender_port": "_1",
+                        "receiver_port": "_1"
+                    },
+                    "Sin_1._2_Cos_2._2": {
+                        "sender": "Sin_1",
+                        "receiver": "Cos_2",
+                        "sender_port": "_2",
+                        "receiver_port": "_2"
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/examples/ONNX/abc_basic-mdf.yaml b/examples/ONNX/abc_basic-mdf.yaml
@@ -0,0 +1,53 @@
+ONNX Model:
+    format: ModECI MDF v0.2
+    generating_application: Python modeci-mdf v0.2.1
+    graphs:
+        torch-jit-export:
+            nodes:
+                Sin_0:
+                    input_ports:
+                        input:
+                            shape: (2, 3)
+                    parameters:
+                        Sin_0:
+                            function: onnx::Sin
+                            args:
+                                input: input
+                    output_ports:
+                        _1:
+                            value: Sin_0
+                Sin_1:
+                    input_ports:
+                        _1:
+                            shape: (2, 3)
+                    parameters:
+                        Sin_1:
+                            function: onnx::Sin
+                            args:
+                                input: _1
+                    output_ports:
+                        _2:
+                            value: Sin_1
+                Cos_2:
+                    input_ports:
+                        _2:
+                            shape: (2, 3)
+                    parameters:
+                        Cos_2:
+                            function: onnx::Cos
+                            args:
+                                input: _2
+                    output_ports:
+                        _3:
+                            value: Cos_2
+            edges:
+                Sin_0._1_Sin_1._1:
+                    sender: Sin_0
+                    receiver: Sin_1
+                    sender_port: _1
+                    receiver_port: _1
+                Sin_1._2_Cos_2._2:
+                    sender: Sin_1
+                    receiver: Cos_2
+                    sender_port: _2
+                    receiver_port: _2
diff --git a/examples/ONNX/abc_basic.py b/examples/ONNX/abc_basic.py
@@ -0,0 +1,84 @@
+"""
+This file does three things:
+    - It implements a simple PyTorch model.
+    - Exports in to ONNX using a combination of tracing and scripting
+    - Converts it to MDF
+"""
+import torch
+import onnx
+
+from onnx import helper
+
+from modeci_mdf.interfaces.onnx import onnx_to_mdf, convert_file
+
+
+class A(torch.nn.Module):
+    def forward(self, x):
+        return torch.sin(x)
+
+
+class B(torch.nn.Module):
+    def forward(self, x):
+        return torch.sin(x)
+
+
+class C(torch.nn.Module):
+    def forward(self, x):
+        return torch.cos(x)
+
+
+class ABC(torch.nn.Module):
+    def __init__(self):
+        super(ABC, self).__init__()
+        self.A = A()
+        self.B = B()
+        self.C = C()
+
+    def forward(self, x):
+
+        # Run A
+        y = self.A(x)
+
+        # Run B (loop_count times)
+        y = self.B(y)
+
+        # Run C
+        y = self.C(y)
+
+        return y
+
+def main():
+
+    model = ABC()
+    dummy_input = torch.zeros(2, 3)
+    # loop_count = torch.tensor(5, dtype=torch.long)
+    torch.onnx.export(model,
+                      (dummy_input),
+                      'abc_basic.onnx',
+                      verbose=True,
+                      input_names=['input'])
+
+
+    # Load it back in using ONNX package
+    onnx_model = onnx.load("abc_basic.onnx")
+    print(onnx_model)
+    onnx.checker.check_model(onnx_model)
+
+    # Extract the loop or if body as a sub-model, this is just because I want
+    # to view it in netron and sub-graphs can't be rendered
+    for node in [node for node in onnx_model.graph.node if node.op_type in ["Loop", 'If']]:
+
+        # Get the GraphProto of the body
+        body_graph = node.attribute[0].g
+
+        # Turn it into a model
+        model_def = helper.make_model(body_graph, producer_name='abc_basic.py')
+
+        onnx.save(model_def, f'examples/{node.name}_body.onnx')
+
+
+    convert_file("abc_basic.onnx")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/src/modeci_mdf/interfaces/onnx/__init__.py b/src/modeci_mdf/interfaces/onnx/__init__.py
@@ -1 +1,3 @@
 from .exporter import onnx_to_mdf, find_subgraphs, convert_file
+
+from .importer import mdf_to_onnx
diff --git a/src/modeci_mdf/interfaces/onnx/exporter.py b/src/modeci_mdf/interfaces/onnx/exporter.py
@@ -121,7 +121,11 @@ def onnx_node_to_mdf(
 
         for p in params_dict:
             if type(params_dict[p]) == Graph:
-                mdf_node.parameters.append(Parameter(id=p, value={'graph_%s'%params_dict[p].id:params_dict[p]}))
+                mdf_node.parameters.append(
+                    Parameter(
+                        id=p, value={"graph_%s" % params_dict[p].id: params_dict[p]}
+                    )
+                )
             else:
                 mdf_node.parameters.append(Parameter(id=p, value=params_dict[p]))
 
@@ -341,22 +345,7 @@ def convert_file(input_file: str):
     onnx.checker.check_model(onnx_model)
     mdf_model = onnx_to_mdf(onnx_model)
     mdf_model.to_json_file(f"{out_filename}.json")
-
-    # Lets convert to YAML as well
-    try:
-        import json
-        import yaml
-
-        with open(f"{out_filename}.yaml", "w") as file:
-            yaml.dump(
-                json.loads(mdf_model.to_json()),
-                file,
-                default_flow_style=None,
-                width=120,
-            )
-        print(f"YAML version written to {out_filename}.yaml")
-    except ImportError as ex:
-        print("Couldn't load pyaml, skipping YAML output.")
+    mdf_model.to_yaml_file(f"{out_filename}.yaml")
 
 
 def main():