Support reshape before gemm

JDAI-CV · Nov 10, 2019 · c7c41be · c7c41be
1 parent 2d26ee0
commit c7c41be
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/tools/onnx2bnn/OnnxConverter.cpp b/tools/onnx2bnn/OnnxConverter.cpp
@@ -260,12 +260,14 @@ std::vector<std::string> OnnxConverter::Convert(
     vector<string> binary_conv_outputs;
     bool has_reshape = false;
     for (const auto &node : model_proto_.graph().node()) {
-        if (has_reshape) {
-            throw std::invalid_argument(
-                "Reshape can only be the last layer for now");
-        }
         NodeAttrHelper helper(node);
         const auto &op = node.op_type();
+        if (has_reshape && op != "Gemm") {
+            throw std::invalid_argument(
+                "Reshape can only be the last layer or precede a gemm layer "
+                "for now");
+        }
+        has_reshape = false;
         VLOG(5) << "Node " << node.name();
         if (op == "Conv") {
             VLOG(5) << "Start converting Conv";