From 7a0c7ef45ca4d46ef4faf4fcfa11820751326878 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 15 Jun 2020 14:12:09 +0800
Subject: [PATCH] feat(mge/module): add module for extern-c-opr

GitOrigin-RevId: a2d9fa067a5db245a3f1546268888a747dab01f1
---
 sdk/c-opr-loaders/mace/dump_model.py | 32 +++++++++++-----------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/sdk/c-opr-loaders/mace/dump_model.py b/sdk/c-opr-loaders/mace/dump_model.py
index ed5e0bb1..21e41022 100644
--- a/sdk/c-opr-loaders/mace/dump_model.py
+++ b/sdk/c-opr-loaders/mace/dump_model.py
@@ -8,9 +8,10 @@
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 import argparse
 
-import megengine._internal as mgb
 import numpy as np
 import yaml
+from megengine import jit
+from megengine.module.external import ExternOprSubgraph
 
 
 # "1,3,224,224" -> (1,3,224,224)
@@ -89,26 +90,19 @@ def main():
             + raw_param
         )
 
-        # cn not ensured
-        cn = mgb.comp_node("xpux")
-        cg = mgb.comp_graph()
-
-        inp = [
-            mgb.make_shared(
-                comp_node=cn,
-                comp_graph=cg,
-                shape=isizes[i],
-                name=input_names[i],
-                dtype=np.float32,
-            )
-            for i in range(len(isizes))
-        ]
+        net = ExternOprSubgraph(wk_raw_content, "mace", osizes)
+        net.eval()
 
-        oup = mgb.opr.extern_c_opr_placeholder(
-            inp, osizes, dump_name="mace", dump_data=wk_raw_content,
-        )
+        @jit.trace(symbolic=True)
+        def inference(inputs):
+            return net(inputs)
+
+        inputs = [
+            np.random.random(isizes[i]).astype(np.float32) for i in range(len(isizes))
+        ]
 
-        mgb.serialize_comp_graph_to_file(args.output, oup)
+        inference.trace(inputs)
+        inference.dump(args.output)
 
 
 if __name__ == "__main__":
-- 
GitLab