diff --git a/imperative/python/megengine/tools/load_network_and_run.py b/imperative/python/megengine/tools/load_network_and_run.py index d226ad1aef49c75a1858c561fde348ed1d59e210..1012e06fccfce558a494b610ef72ba56994eff85 100755 --- a/imperative/python/megengine/tools/load_network_and_run.py +++ b/imperative/python/megengine/tools/load_network_and_run.py @@ -121,6 +121,9 @@ def run_model(args, graph, inputs, outputs, data): # must use level0 to avoid unintended opr modification graph.options.graph_opt_level = 0 + if args.weight_preprocess: + graph.enable_weight_preprocess() + logger.info("input tensors: ") for k, v in data.items(): logger.info(" {}: {}".format(k, v.shape)) @@ -161,8 +164,8 @@ def run_model(args, graph, inputs, outputs, data): func.wait() return [oup_node.get_value().numpy() for oup_node in output_dict.values()] - if args.warm_up: - logger.info("warming up") + for i in range(args.warm_up): + logger.info("warming up {}".format(i)) run() total_time = 0 @@ -276,8 +279,9 @@ def main(): ) parser.add_argument( "--warm-up", - action="store_true", - help="warm up model before do timing " " for better estimation", + type=int, + default=0, + help="times of warm up model before do timing " " for better estimation", ) parser.add_argument( "--verbose", @@ -394,6 +398,13 @@ def main(): parser.add_argument( "--custom-op-lib", type=str, help="path of the custom op", ) + parser.add_argument( + "--weight-preprocess", + action="store_true", + help="Execute operators with weight preprocess, which can" + "optimize the operator execution time with algo of winograd," + "im2col ,etc.,but it may consume more memory.", + ) args = parser.parse_args() diff --git a/imperative/python/src/graph_rt.cpp b/imperative/python/src/graph_rt.cpp index de7f28b99ae6da0c4f7143407af9a3c9ee80364a..a6085aef1b405939551452c4c879e44e20901307 100644 --- a/imperative/python/src/graph_rt.cpp +++ b/imperative/python/src/graph_rt.cpp @@ -253,6 +253,10 @@ void init_graph_rt(py::module m) { } return graph.compile(spec); }) + .def("enable_weight_preprocess", + [](cg::ComputingGraph& graph) { + graph.options().graph_opt.enable_weight_preprocess(); + }) .def_property_readonly( "options", py::overload_cast<>(&cg::ComputingGraph::options));