fix duplicate layer names bug (#446) (#467)

Suppose we have a network with (not all distinct) layer names

layer
layer_1
layer

When ImporterContext sees "layer", it sees it's not in mLayerNameCounts,
and sets mLayerNameCounts["layer"] = 1 and adds a TRT layer with name "layer".
It then sees "layer_1", concludes it's not in mLayerNameCounts, so it sets
mLayerNameCounts["layer_1"] = 1 and adds a TRT layer with name "layer_1".
NOW when it sees "layer", it sees that mLayerNameCounts["layer"] == 1,
so we produce a "uniqueName" of
"layer" + "_" + std::to_string(mLayerNameCounts["layer"] ), ie "layer_1",
which is a name conflict for the TRT net.

This change keeps track of all inserted names in a set and in the case of
duplicates, tries suffix-appended modifications of the duplicated name
by ever increasing integers until a name appears which has not been used.

Author: tdp2110

ImporterContext.hpp

@@ -26,11 +26,12 @@
 #include "onnx2trt_utils.hpp"
 
 #include <list>
+#include <set>
+#include <string>
 #include <unordered_map>
 
 namespace onnx2trt
 {
-
 class ImporterContext final : public IImporterContext
 {
     nvinfer1::INetworkDefinition* _network;
@@ -44,13 +45,17 @@ class ImporterContext final : public IImporterContext
     StringMap<float> mTensorRangeMins;
     StringMap<float> mTensorRangeMaxes;
     StringMap<nvinfer1::DataType> mLayerPrecisions;
-    StringMap<size_t>
-        mTensorNameCounts; // Keep track of how many times a tensor name shows up, to avoid duplicate naming in TRT.
-    StringMap<size_t>
-        mLayerNameCounts; // Keep track of how many times a tensor name shows up, to avoid duplicate naming in TRT.
-    std::unordered_set<std::string> mUnsupportedShapeTensors; // Container to hold output tensor names of layers that produce shape tensor outputs but do not natively support them.
-    StringMap<std::string> mLoopTensors; // Container to map subgraph tensors to their original outer graph names.
-    std::string mOnnxFileLocation; // Keep track of the directory of the parsed ONNX file
+    std::set<std::string> mTensorNames;                       // keep track of tensor names used so far,
+                                                              // to avoid duplicate naming in TRT.
+    std::set<std::string> mLayerNames;                        // keep track of layer names used so far,
+                                                              // to avoid duplicate naming in TRT.
+    int64_t mSuffixCounter = 0;                               // increasing suffix counter used to uniquify layer names.
+    std::unordered_set<std::string> mUnsupportedShapeTensors; // Container to hold any shape tensors that are
+                                                              // the output of layers that do not support
+                                                              // shape tensors.
+    StringMap<std::string> mLoopTensors;                      // Container to map subgraph tensors to
+                                                              // their original outer graph names.
+    std::string mOnnxFileLocation;                            // Keep track of the directory of the parsed ONNX file
 
 public:
     ImporterContext(nvinfer1::INetworkDefinition* network, nvinfer1::ILogger* logger)
@@ -98,13 +103,12 @@ class ImporterContext final : public IImporterContext
     {
         return mOnnxFileLocation;
     }
-    // This actually handles weights as well, but is named this way to be consistent with the tensors()
+    // This actually handles weights as well, but is named this way to be
+    // consistent with the tensors()
     virtual void registerTensor(TensorOrWeights tensor, const std::string& basename) override
     {
         // TRT requires unique tensor names.
-        const std::string uniqueName
-            = mTensorNameCounts[basename] ? (basename + "_" + std::to_string(mTensorNameCounts[basename])) : basename;
-        ++mTensorNameCounts[basename];
+        const std::string uniqueName = generateUniqueName(mTensorNames, basename);
 
         if (tensor)
         {
@@ -122,8 +126,9 @@ class ImporterContext final : public IImporterContext
                     convertINT64(reinterpret_cast<int64_t*>(weights.values), weights.shape, ctx), weights.shape};
             }
         }
-        // Overwrite previous tensors registered with the same name (this only happens when there are subgraphs,
-        // and in that case, overwriting is the desired behavior).
+        // Overwrite previous tensors registered with the same name (this only
+        // happens when there are subgraphs, and in that case, overwriting is the
+        // desired behavior).
         this->tensors()[basename] = std::move(tensor);
     }
 
@@ -133,9 +138,7 @@ class ImporterContext final : public IImporterContext
         if (layer)
         {
             const std::string name = basename.empty() ? layer->getName() : basename;
-            const std::string uniqueName
-                = mLayerNameCounts[name] ? (name + "_" + std::to_string(mLayerNameCounts[name])) : name;
-            ++mLayerNameCounts[name];
+            const std::string uniqueName = generateUniqueName(mLayerNames, basename);
 
             auto* ctx = this; // To enable logging.
             LOG_VERBOSE("Registering layer: " << name << " for ONNX node: " << basename);
@@ -225,6 +228,22 @@ class ImporterContext final : public IImporterContext
             return _opsets.at(domain);
         }
     }
+
+private:
+    std::string generateUniqueName(std::set<std::string>& namesSet, const std::string& basename)
+    {
+        std::string candidate = basename;
+
+        while (namesSet.find(candidate) != namesSet.end())
+        {
+            candidate = basename + "_" + std::to_string(mSuffixCounter);
+            ++mSuffixCounter;
+        }
+
+        namesSet.insert(candidate);
+
+        return candidate;
+    }
 };
 
 } // namespace onnx2trt