Enable Pruning transformation by default inside Model Optimizer (open…

…vinotoolkit#5633)
tomdol · May 18, 2021 · 8be740a · 8be740a
1 parent b835c6e
commit 8be740a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/inference-engine/src/offline_transformations/src/moc_transformations.cpp b/inference-engine/src/offline_transformations/src/moc_transformations.cpp
@@ -5,10 +5,16 @@
 #include <memory>
 
 #include "moc_transformations.hpp"
+#include "pruning.hpp"
 
+#include <ngraph/pass/manager.hpp>
 
 NGRAPH_RTTI_DEFINITION(ngraph::pass::MOCTransformations, "MOCTransformations", 0);
 
-bool ngraph::pass::MOCTransformations::run_on_function(std::shared_ptr<ngraph::Function>) {
+bool ngraph::pass::MOCTransformations::run_on_function(std::shared_ptr<ngraph::Function> f) {
+    ngraph::pass::Manager m(get_pass_config());
+    m.register_pass<Pruning>();
+    m.run_passes(f);
+
     return false;
 }
diff --git a/model-optimizer/mo/back/offline_transformations.py b/model-optimizer/mo/back/offline_transformations.py
@@ -35,6 +35,7 @@ def apply_offline_transformations(input_model: str, framework: str, transforms:
 
         available_transformations[name](net, **args)
 
+    ApplyMOCTransformations(net, False)
     net.serialize(input_model + ".xml", input_model + ".bin")
     path_to_mapping = input_model + ".mapping"
     GenerateMappingFile(net, path_to_mapping.encode('utf-8'), extract_names)