make strkeymap has const, check prelu slope

JDAI-CV · Jul 16, 2019 · b1a25f6 · b1a25f6
1 parent 51f2779
commit b1a25f6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/include/common/StrKeyMap.h b/include/common/StrKeyMap.h
@@ -40,7 +40,7 @@ class StrKeyMap {
     void clear() {
         map_.clear();
     }
-    bool has(const std::string &key) {
+    bool has(const std::string &key) const {
         return map_.find(key) != map_.end();
     }
 

diff --git a/tools/onnx2daq/OnnxConverter.cpp b/tools/onnx2daq/OnnxConverter.cpp
@@ -1078,9 +1078,13 @@ std::pair<bool, std::string> OnnxConverter::IsNodeSupported(
         }
     } else if (op == "PRelu") {
         const auto slope_name = m(node.input(1));
-        if (onnx_tensors_.at(slope_name).shape != Shape{1}) {
-            // TODO: support it
-            return {false, "Only support one element slope."};
+        if (onnx_tensors_.has(slope_name)) {
+            if (onnx_tensors_.at(slope_name).shape != Shape{1}) {
+                // TODO: support it
+                return {false, "PRelu only support one element slope."};
+            }
+        } else {
+            return {false, "PRelu slope must be known"};
         }
     } else if (op == "Gemm") {
         const auto transA = helper.get("transA", 0);