[mir_caffe] Inherit tensor names (#6845)

Set `Operation::Output`s names based on model tensor names.

Signed-off-by: Sergei Barannikov <s.barannikov@samsung.com>

diff --git a/compiler/mir-caffe-importer/caffe_importer.cpp b/compiler/mir-caffe-importer/caffe_importer.cpp
index 359a97d..ade1ea2 100644 (file)
--- a/compiler/mir-caffe-importer/caffe_importer.cpp
+++ b/compiler/mir-caffe-importer/caffe_importer.cpp
@@ -178,9 +178,9 @@ void CaffeImporter::createMIRNodesFromLayer(const LayerParameter &layer)
       assert(false && "All unsupported types should have been found before this pass.");
   }
 
-  assert(layer.top_size() == static_cast<int>(outputs.size()) && "Number of outputs differs.");
+  assert(static_cast<int>(outputs.size()) == layer.top_size() && "Number of outputs differs.");
   for (int i = 0; i < layer.top_size(); ++i)
-    _blobNameToOpOutput[layer.top(i)] = outputs.at(i);
+    setOutputForBlob(layer.top(i), outputs[i]);
 }
 
 void CaffeImporter::collectUnsupportedOp(const LayerParameter &lp)
@@ -246,18 +246,36 @@ CaffeImporter::getMIRInputsForLayer(const LayerParameter &layer)
   std::vector<mir::Operation::Output *> inputs;
 
   for (const auto &input_name : layer.bottom())
-    inputs.push_back(_blobNameToOpOutput.at(input_name));
+    inputs.push_back(getOutputForBlob(input_name));
 
   return inputs;
 }
 
+mir::Operation::Output *CaffeImporter::getOutputForBlob(const std::string &blob_name) const
+{
+  return _blobNameToOpOutput.at(blob_name);
+}
+
+void CaffeImporter::setOutputForBlob(const std::string &blob_name, mir::Operation::Output *output)
+{
+  const auto it = _blobNameToOpOutput.find(blob_name);
+  if (it != _blobNameToOpOutput.cend())
+  {
+    // caffe input blob name could be same as output blob name, and next line will overwrite
+    // '_blobNameToOpOutput' element, but in all networks that I saw it was not a problem
+    it->second->setName("");
+  }
+  output->setName(blob_name);
+  _blobNameToOpOutput[blob_name] = output;
+}
+
 void CaffeImporter::setGraphOutputs()
 {
   const auto &last_layer = _net->layer(_net->layer_size() - 1);
-  // For now, we assume that:
+  // TODO For now, we assume that:
   //   - there is exactly one output;
   //   - the output is from the last layer.
-  auto output = _blobNameToOpOutput[last_layer.top(0)];
+  auto output = getOutputForBlob(last_layer.top(0));
   _graph->create<mir::ops::OutputOp>(output->getNode()->getName(), output);
   output->getNode()->setName("");
 }

diff --git a/compiler/mir-caffe-importer/caffe_importer.h b/compiler/mir-caffe-importer/caffe_importer.h
index 4726c7b..4fb609f 100644 (file)
--- a/compiler/mir-caffe-importer/caffe_importer.h
+++ b/compiler/mir-caffe-importer/caffe_importer.h
@@ -70,6 +70,9 @@ private:
    */
   void createMIRNodesFromLayer(const ::caffe::LayerParameter &layer);
 
+  mir::Operation::Output *getOutputForBlob(const std::string &blob_name) const;
+  void setOutputForBlob(const std::string &blob_name, mir::Operation::Output *output);
+
   /**
    * @brief Collect unsupported parts of caffe layer
    */