Merge pull request #1167 from pytorch/fix_renaming_itensor

fix: converter renaming already named tensors

Author: peri044

core/conversion/conversion.cpp

@@ -188,7 +188,7 @@ void AddInputs(
       ctx->input_is_dynamic = true;
     }
 
-    ctx->value_tensor_map[in] = trt_in;
+    ctx->RecordNewITensor(in, trt_in);
     ctx->num_inputs += 1;
   }
 

core/conversion/conversionctx/ConversionCtx.cpp

@@ -143,8 +143,8 @@ ConversionCtx::~ConversionCtx() {
 }
 
 nvinfer1::ITensor* ConversionCtx::AssociateValueAndTensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor) {
-  tensor->setName(value->debugName().c_str());
-  this->value_tensor_map[value] = tensor;
+  RecordNewITensor(value, tensor);
+
   return tensor;
 }
 
@@ -153,6 +153,15 @@ torch::jit::IValue* ConversionCtx::AssociateValueAndIValue(const torch::jit::Val
   return &this->evaluated_value_map[value];
 }
 
+void ConversionCtx::RecordNewITensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor) {
+  value_tensor_map[value] = tensor;
+  auto ret = seen_itensors.insert(tensor);
+  if (!ret.second) {
+    LOG_WARNING(
+        "Trying to record the value " << value->debugName() << " with the ITensor " << tensor->getName() << " again.");
+  }
+}
+
 std::string ConversionCtx::SerializeEngine() {
 #if NV_TENSORRT_MAJOR > 7
   auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);

core/conversion/conversionctx/ConversionCtx.h

@@ -48,6 +48,7 @@ struct ConversionCtx {
   ConversionCtx(BuilderSettings settings);
   std::string SerializeEngine();
   nvinfer1::ITensor* AssociateValueAndTensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor);
+  void RecordNewITensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor);
   torch::jit::IValue* AssociateValueAndIValue(const torch::jit::Value* value, torch::jit::IValue tensor);
   bool CheckLayerAddition(const torch::jit::Node* n);
 
@@ -71,6 +72,9 @@ struct ConversionCtx {
 
   std::unordered_map<const torch::jit::Value*, nvinfer1::ITensor*> value_tensor_map;
   std::unordered_map<const torch::jit::Value*, torch::jit::IValue> evaluated_value_map;
+
+  // record already named ITensors to prevent rewriting another name to the same tensor
+  std::unordered_set<nvinfer1::ITensor*> seen_itensors;
 };
 
 } // namespace conversion