diff --git a/core/conversion/conversion.cpp b/core/conversion/conversion.cpp index 0fa805b7fe..3211e7dd98 100644 --- a/core/conversion/conversion.cpp +++ b/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; } diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp index 688aaa7939..a24a15904c 100644 --- a/core/conversion/conversionctx/ConversionCtx.cpp +++ b/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); diff --git a/core/conversion/conversionctx/ConversionCtx.h b/core/conversion/conversionctx/ConversionCtx.h index 3744ff6213..988a58bf49 100644 --- a/core/conversion/conversionctx/ConversionCtx.h +++ b/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 value_tensor_map; std::unordered_map evaluated_value_map; + + // record already named ITensors to prevent rewriting another name to the same tensor + std::unordered_set seen_itensors; }; } // namespace conversion