Add extremesoft level

JDAI-CV · May 15, 2019 · 0b3e8e5 · 0b3e8e5
1 parent d55df70
commit 0b3e8e5
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 8 deletions.
diff --git a/third_party/onnx b/third_party/onnx
diff --git a/tools/onnx2bnn/OnnxConverter.cpp b/tools/onnx2bnn/OnnxConverter.cpp
@@ -213,7 +213,7 @@ vector<bin_t> bitpack(const float *data, Shape shape) {
 
 void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto &model_proto,
                             const std::string &filepath,
-                            const bool strict) {
+                            const OnnxConverter::Level level) {
     GOOGLE_PROTOBUF_VERIFY_VERSION;
 
     // We recognize binary convolutions in our custom ONNX optimizers.
@@ -223,9 +223,12 @@ void OnnxConverter::Convert(const ONNX_NAMESPACE::ModelProto &model_proto,
     // https://github.com/daquexian/onnx/blob/optimizer_for_bnn/onnx/optimizer/passes/dabnn_bconv_strict.h
     // for details.
     vector<string> optimizers{"eliminate_nop_pad", "dabnn_bconv_strict"};
-    if (!strict) {
+    if (level == Level::kSoft || level == Level::kExtremeSoft) {
         optimizers.push_back("dabnn_bconv_soft");
     }
+    if (level == Level::kExtremeSoft) {
+        optimizers.push_back("dabnn_bconv_extreme_soft");
+    }
     // model_proto is only used here. Please use the member variable model_proto_
     // in the following code
     model_proto_ = ONNX_NAMESPACE::optimization::Optimize(

diff --git a/tools/onnx2bnn/OnnxConverter.h b/tools/onnx2bnn/OnnxConverter.h
@@ -148,9 +148,14 @@ class OnnxConverter {
     }
 
    public:
+    enum class Level {
+        kStrict,
+        kSoft,
+        kExtremeSoft,
+    };
     void Convert(const ONNX_NAMESPACE::ModelProto &model,
                  const std::string &filepath,
-                 const bool strict=false);
+                 const Level level=Level::kSoft);
 };
 
 template <>

diff --git a/tools/onnx2bnn/onnx2bnn.cpp b/tools/onnx2bnn/onnx2bnn.cpp
@@ -20,13 +20,21 @@ void usage(const std::string &filename) {
 }
 
 int main(int argc, char **argv) {
-    argh::parser cmdl(argv);
+    argh::parser cmdl(argc, argv);
     google::InitGoogleLogging(cmdl[0].c_str());
     FLAGS_alsologtostderr = true;
-    if (argc < 3) {
-        usage(argv[0]);
+    if (!cmdl(2)) {
+        usage(cmdl[0]);
         return -1;
     }
+    bnn::OnnxConverter::Level opt_level = bnn::OnnxConverter::Level::kSoft;
+    if (cmdl["strict"]) {
+        opt_level = bnn::OnnxConverter::Level::kStrict;
+    }
+    if (cmdl["extremesoft"]) {
+        opt_level = bnn::OnnxConverter::Level::kExtremeSoft;
+    }
+
     ONNX_NAMESPACE::ModelProto model_proto;
     {
         std::ifstream ifs(cmdl[1], std::ios::in | std::ios::binary);
@@ -35,7 +43,7 @@ int main(int argc, char **argv) {
     }
 
     bnn::OnnxConverter converter;
-    converter.Convert(model_proto, cmdl[2], cmdl["strict"]);
+    converter.Convert(model_proto, cmdl[2], opt_level);
 
     google::protobuf::ShutdownProtobufLibrary();
     return 0;