[IE] Blob dumper from IR #2

eshoguli · Mar 1, 2021 · efca024 · efca024
1 parent 9b52c90
commit efca024
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 1 deletion.
diff --git a/inference-engine/samples/blob_dumper/blob_dumper.hpp b/inference-engine/samples/blob_dumper/blob_dumper.hpp
@@ -21,6 +21,12 @@ static const char dump_dir_message[] = "Required. Directory to store dumps.";
 /// @brief message for reduce argument
 static const char reduce_message[] = "Optional. Reduce blobs.";
 
+/// @brief message for rename_message argument
+static const char rename_message[] = "Optional. Rename operation names.";
+
+/// @brief message for ignore_const_message argument
+static const char ignore_const_message[] = "Optional. Ignore Const operations.";
+
 /// @brief Define flag for showing help message <br>
 DEFINE_bool(h, false, help_message);
 
@@ -35,10 +41,17 @@ DEFINE_string(model, "", model_message);
 /// It is a required parameter
 DEFINE_string(dump_dir, "", dump_dir_message);
 
-/// @brief Define parameter for set model file <br>
+/// @brief Define parameter to reduce blobs <br>
 /// It is a optional parameter
 DEFINE_bool(reduce, false, reduce_message);
 
+/// @brief Define parameter to rename operation names <br>
+/// It is a optional parameter
+DEFINE_bool(rename, false, rename_message);
+
+/// @brief Define parameter to ignore constant operation dump<br>
+/// It is a optional parameter
+DEFINE_bool(ignore_const, false, ignore_const_message);
 
 /**
 * @brief This function show a help message

diff --git a/inference-engine/samples/blob_dumper/main.cpp b/inference-engine/samples/blob_dumper/main.cpp
@@ -26,6 +26,39 @@
 
 using namespace InferenceEngine;
 
+class OperationName {
+public:
+    enum sortType {
+        byName,
+        byInference
+    };
+
+    OperationName(const OperationName::sortType sort) : sort(sort), executionIndex(0ul) {}
+
+    std::string getName(const std::string& operationType) {
+        const auto it = operationIndexes.find(operationType);
+
+        size_t index = 0ul;
+        if (it == operationIndexes.end()) {
+            operationIndexes[operationType] = 0;
+            index = 0ul;
+        }
+        else {
+            index = it->second;
+            operationIndexes[operationType] = ++index;
+        }
+
+        return (sort == OperationName::sortType::byName) ?
+            operationType + "_" + std::to_string(index) :
+            std::to_string(executionIndex++) + "_" + operationType + "_" + std::to_string(index);
+    }
+
+private:
+    std::map<std::string, size_t> operationIndexes;
+    size_t executionIndex;
+    OperationName::sortType sort;
+};
+
 bool ParseAndCheckCommandLine(int argc, char *argv[]) {
     // ---------------------------Parsing and validating input arguments--------------------------------------
     slog::info << "Parsing input parameters" << slog::endl;
@@ -200,7 +233,13 @@ int main(int argc, char *argv[]) {
             THROW_IE_EXCEPTION << "not found";
         };
 
+        OperationName operationName(OperationName::byInference);
+
         for (auto layer : layers) {
+            if (FLAGS_ignore_const && (layer->type == "Const")) {
+                continue;
+            }
+
             if (!layer->blobs.empty()) {
                 for (auto blobIt : layer->blobs) {
                     CNNLayerPtr childLayerToBuildName;
@@ -280,6 +319,10 @@ int main(int argc, char *argv[]) {
                         }
                     }
 
+                    if (FLAGS_rename) {
+                        nodeName = operationName.getName(layer->type);
+                    }
+
                     const std::string blobName = blobIt.first;
                     const std::string dumpFilePath = FLAGS_dump_dir + "\\" + nodeName + "_" + blobName + ".ieb";
 
@@ -304,6 +347,10 @@ int main(int argc, char *argv[]) {
                 }
                 nodeName += ss.str();
 
+                if (FLAGS_rename) {
+                    nodeName = operationName.getName(layer->type);
+                }
+
                 const std::string dumpFilePath = FLAGS_dump_dir + "\\" + nodeName + ".ieb";
 
                 BlobDumper dumper(nullptr);