feat(//core/conversion/conversionctx): Move inline function to associate

narendasan · narendasan · commit 6ab9814e4297 · 2020-03-27T01:02:09.000-07:00
values and tensors to a method in the conversion ctx

Signed-off-by: Naren Dasan &lt;naren@narendasan.com&gt;
Signed-off-by: Naren Dasan &lt;narens@nvidia.com&gt;
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -37,11 +37,11 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
     switch(settings.op_precision) {
     case nvinfer1::DataType::kHALF:
         cfg->setFlag(nvinfer1::BuilderFlag::kFP16);
-        input_type = nvinfer1::DataType::kHALF; 
+        input_type = nvinfer1::DataType::kHALF;
         break;
     // case nvinfer1::DataType::kINT8:
     //     cfg->setFlag(nvinfer1::BuilderFlag::kINT8);
-    //         input_type = nvinfer1::DataType::kFLOAT;    
+    //         input_type = nvinfer1::DataType::kFLOAT;
     //     break;
     case nvinfer1::DataType::kFLOAT:
     default:
@@ -80,13 +80,30 @@ ConversionCtx::~ConversionCtx() {
         free(ptr);
     }
 }
-    
+
+nvinfer1::ITensor* ConversionCtx::AssociateValueAndTensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor) {
+    tensor->setName(value->debugName().c_str());
+    this->value_tensor_map[value] = tensor;
+    return tensor;
+}
+
 std::string ConversionCtx::SerializeEngine() {
     auto engine = builder->buildEngineWithConfig(*net, *cfg);
     auto serialized_engine = engine->serialize();
     return std::string((const char*)serialized_engine->data(), serialized_engine->size());
 }
 
+bool ConversionCtx::CheckLayerAddition(const torch::jit::Node* n) {
+    for (auto out : n->outputs()) {
+        auto iter = this->value_tensor_map.find(out);
+        if (iter == this->value_tensor_map.end()) {
+            LOG_WARNING("Node " << util::node_info(n) << " output: " << out->debugName() << " does not have a coresponding output, may potentially indicate a defective converter");
+            return false;
+        }
+    }
+    return true;
+}
+
 } // namespace conversion
 } // namespace core
 } // namespace trtorch
diff --git a/core/conversion/conversionctx/ConversionCtx.h b/core/conversion/conversionctx/ConversionCtx.h
@@ -30,12 +30,15 @@ struct BuilderSettings {
 
     BuilderSettings() = default;
     BuilderSettings(const BuilderSettings& other) = default;
-    friend std::ostream& operator<<(std::ostream& os, const BuilderSettings& s); 
+    friend std::ostream& operator<<(std::ostream& os, const BuilderSettings& s);
 };
-    
+
 struct ConversionCtx {
     ConversionCtx(BuilderSettings settings);
     std::string SerializeEngine();
+    nvinfer1::ITensor* AssociateValueAndTensor(const torch::jit::Value* value, nvinfer1::ITensor* tensor);
+    bool CheckLayerAddition(const torch::jit::Node* n);
+
     ~ConversionCtx();
 
     nvinfer1::IBuilder* builder;
@@ -50,12 +53,12 @@ struct ConversionCtx {
     // is constructed from a PyTorch Tensor it allocates the data here to store a
     // copy of the values
     std::vector<void*> builder_resources;
-    
+
     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;
 };
 
-} // namespace conversion    
+} // namespace conversion
 } // namespace core
 } // namespace trtorch
-    
+
diff --git a/core/conversion/converters/converters.h b/core/conversion/converters/converters.h
@@ -69,12 +69,12 @@ class Arg {
     ArgContainer ptr_;
     Type type_;
 };
-        
-    
+
+
 
 typedef std::vector<Arg> args;
 typedef std::function<bool(ConversionCtx*, const torch::jit::Node*, args&)> OpConverter;
-struct ConversionPattern {    
+struct ConversionPattern {
     std::string signature;
     OpConverter converter;
 };
@@ -107,20 +107,14 @@ struct Weights {
     Weights();
     Weights(ConversionCtx* ctx, at::Tensor t);
     Weights(ConversionCtx* ctx, float val);
-    friend std::ostream& operator<<(std::ostream& os, const Weights& w); 
+    friend std::ostream& operator<<(std::ostream& os, const Weights& w);
 };
 
 inline nvinfer1::ITensor* tensor_to_const(ConversionCtx* ctx, at::Tensor t) {
     auto t_weights = Weights(ctx, t);
     return ctx->net->addConstant(t_weights.shape, t_weights.data)->getOutput(0);
 }
 
-inline nvinfer1::ITensor* associate_value_and_tensor(ConversionCtx* ctx, const torch::jit::Value* value, nvinfer1::ITensor* tensor) {
-    tensor->setName(value->debugName().c_str());
-    ctx->value_tensor_map[value] = tensor;
-    return tensor;
-}
-
 } // namespace converters
 } // namespace conversion
 } // namespace core
diff --git a/core/conversion/converters/impl/activation.cpp b/core/conversion/converters/impl/activation.cpp
@@ -18,7 +18,7 @@ namespace {
                   "Unable to create " #act " layer from node: " << *n);        \
                                                                                \
     new_layer->setName(util::node_info(n).c_str());                            \
-    associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0)); \
+    ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));    \
     LOG_DEBUG("Output tensor shape: "                                          \
                << new_layer->getOutput(0)->getDimensions());                   \
                                                                                \
diff --git a/core/conversion/converters/impl/batch_norm.cpp b/core/conversion/converters/impl/batch_norm.cpp
@@ -32,10 +32,10 @@ bool ConvertConvBatchNorm(ConversionCtx* ctx, const torch::jit::Node* n, args& a
 
     auto bn_as_conv = ctx->net->addConvolutionNd(*input, weights.num_output_maps, weights.kernel_shape, weights.data, bias.data);
     TRTORCH_CHECK(bn_as_conv, "Unable to create fused batch norm from node: " << *n);
-    
+
     bn_as_conv->setName(util::node_info(n).c_str());
 
-    auto bn_out = associate_value_and_tensor(ctx, n->outputs()[0], bn_as_conv->getOutput(0));
+    auto bn_out = ctx->AssociateValueAndTensor(n->outputs()[0], bn_as_conv->getOutput(0));
     LOG_DEBUG("Output tensor shape: " << bn_out->getDimensions());
     return true;
 }
@@ -67,25 +67,25 @@ bool ConvertLinearBatchNorm(ConversionCtx* ctx, const torch::jit::Node* n, args&
     auto bn_biased_out = bn_biased->getOutput(0);
 
     bn_biased->setName(util::node_info(n).c_str());
-    associate_value_and_tensor(ctx, n->outputs()[0], bn_biased_out);
+    ctx->AssociateValueAndTensor(n->outputs()[0], bn_biased_out);
 
     return true;
 }
 
 volatile auto batch_norm_registrations = RegisterNodeConversionPatterns()
     .pattern({
-            R"SIG(aten::batch_norm(Tensor input, Tensor? gamma, Tensor? beta, 
-                               Tensor? mean, Tensor? var, 
+            R"SIG(aten::batch_norm(Tensor input, Tensor? gamma, Tensor? beta,
+                               Tensor? mean, Tensor? var,
                                bool training, float momentum, float eps, bool cudnn_enabled) -> (Tensor))SIG",
             [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
                 auto input = args[0].ITensor();
                 auto shape = input->getDimensions();
                 auto gamma = args[1].unwrapToTensor();
-                
+
                 if (/*training*/ args[5].unwrapToBool()) {
                     LOG_WARNING("TensorRT only converts forward pass of graphs, but saw training = True, may see undefined behavior, consider placing module in eval mode");
                 }
-                
+
                 // If gamma is None this fails
                 if (util::volume(shape) == gamma.numel()) {
                     return ConvertLinearBatchNorm(ctx, n, args);
@@ -101,4 +101,4 @@ volatile auto batch_norm_registrations = RegisterNodeConversionPatterns()
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace trtorch 
+} // namespace trtorch
diff --git a/core/conversion/converters/impl/constant.cpp b/core/conversion/converters/impl/constant.cpp
@@ -19,10 +19,10 @@ auto constant_registrations = RegisterNodeConversionPatterns()
             auto t_weights = Weights(ctx, t);
             auto const_layer = ctx->net->addConstant(t_weights.shape, t_weights.data);
             const_layer->setName(util::node_info(n).c_str());
-            auto const_out = associate_value_and_tensor(ctx, n->outputs()[0], const_layer->getOutput(0));
+            auto const_out = ctx->AssociateValueAndTensor(n->outputs()[0], const_layer->getOutput(0));
 
             LOG_DEBUG("Output tensor shape: " << const_out->getDimensions());
-            
+
             return true;
         }
     });
@@ -31,5 +31,5 @@ auto constant_registrations = RegisterNodeConversionPatterns()
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // namespace trtorch 
+} // namespace trtorch
 
diff --git a/core/conversion/converters/impl/conv_deconv.cpp b/core/conversion/converters/impl/conv_deconv.cpp
@@ -9,14 +9,14 @@ namespace impl {
 namespace {
 auto conv_registrations = RegisterNodeConversionPatterns()
     .pattern({
-        R"SIG(aten::_convolution(Tensor input, Tensor weight, 
+        R"SIG(aten::_convolution(Tensor input, Tensor weight,
                                  Tensor? bias, int[] stride, int[] padding,
-                                 int[] dilation, bool transposed, 
-                                 int[] output_padding, int groups, bool benchmark, 
+                                 int[] dilation, bool transposed,
+                                 int[] output_padding, int groups, bool benchmark,
                                  bool deterministic, bool cudnn_enabled) -> (Tensor))SIG",
         [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
             auto in = args[0].ITensor();
-                        
+
             auto w = Weights(ctx, args[1].unwrapToTensor());
             auto stride = util::toDimsHW(args[3].unwrapToIntList());
             LOG_DEBUG("stride: " << stride);
@@ -27,11 +27,11 @@ auto conv_registrations = RegisterNodeConversionPatterns()
             bool transposed = args[6].unwrapToBool();
             auto out_padding = util::toDimsHW(args[7].unwrapToIntList());
             LOG_DEBUG("out_padding: " << out_padding);
-            int64_t groups = args[8].unwrapToInt(); 
-            
+            int64_t groups = args[8].unwrapToInt();
+
             nvinfer1::ILayer* new_layer;
             if (transposed) {
-                //TODO: Check deconv correctness 
+                //TODO: Check deconv correctness
                 LOG_WARNING(ctx->logger, "Deconvolution converter has not be tested");
                 nvinfer1::IDeconvolutionLayer* deconv;
                 if (args[2].IValue()->isTensor()) {
@@ -54,9 +54,9 @@ auto conv_registrations = RegisterNodeConversionPatterns()
                 } else {
                     conv = ctx->net->addConvolutionNd(*in, w.num_output_maps, w.kernel_shape, w.data, Weights().data);
                 }
-                
+
                 TRTORCH_CHECK(conv, "Unable to create convolution layer from node: " << *n);
-                
+
                 conv->setStrideNd(stride);
                 conv->setPaddingMode(nvinfer1::PaddingMode::kCAFFE_ROUND_DOWN);
                 conv->setPaddingNd(padding);
@@ -67,7 +67,7 @@ auto conv_registrations = RegisterNodeConversionPatterns()
             }
             new_layer->setName(util::node_info(n).c_str());
 
-            auto out = associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0));
+            auto out = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
 
             LOG_DEBUG("Output tensor shape: " << out->getDimensions());
 
@@ -79,4 +79,4 @@ auto conv_registrations = RegisterNodeConversionPatterns()
 } // namespace converters
 } // namespace conversion
 } // namespace core
-} // trtorch 
+} // trtorch
diff --git a/core/conversion/converters/impl/linear.cpp b/core/conversion/converters/impl/linear.cpp
@@ -52,7 +52,7 @@ auto linear_registrations = RegisterNodeConversionPatterns()
             TRTORCH_CHECK(new_layer,"Unable to create linear layer from node: " << *n);
 
             new_layer->setName(util::node_info(n).c_str());
-            auto out_tensor = associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0));
+            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
 
             LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
 
diff --git a/core/conversion/converters/impl/pooling.cpp b/core/conversion/converters/impl/pooling.cpp
@@ -52,7 +52,7 @@ auto pooling_registrations = RegisterNodeConversionPatterns()
             new_layer->setPaddingMode(padding_mode);
 
             new_layer->setName(util::node_info(n).c_str());
-            auto out_tensor = associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0));
+            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
 
             LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
             return true;
@@ -97,7 +97,7 @@ auto pooling_registrations = RegisterNodeConversionPatterns()
             new_layer->setStrideNd(util::toDims(stride));
 
             new_layer->setName(util::node_info(n).c_str());
-            auto out_tensor = associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0));
+            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
 
             LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
             return true;
diff --git a/core/conversion/converters/impl/unary.cpp b/core/conversion/converters/impl/unary.cpp
@@ -23,7 +23,7 @@ namespace {
                  "Unable to create " #unary " layer from node: " << *n);        \
                                                                                 \
              unary->setName(util::node_info(n).c_str());                        \
-             auto out_tensor = associate_value_and_tensor(ctx,                  \
+             auto out_tensor = ctx->AssociateValueAndTensor(                    \
                                                           n->outputs()[0],      \
                                                           unary->getOutput(0)); \
              LOG_DEBUG(                                                         \

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ namespace {`
`18`	`18`	`"Unable to create " #act " layer from node: " << *n); \`
`19`	`19`	`\`
`20`	`20`	`new_layer->setName(util::node_info(n).c_str()); \`
`21`		`- associate_value_and_tensor(ctx, n->outputs()[0], new_layer->getOutput(0)); \`
	`21`	`+ ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0)); \`
`22`	`22`	`LOG_DEBUG("Output tensor shape: " \`
`23`	`23`	`<< new_layer->getOutput(0)->getDimensions()); \`
`24`	`24`	`\`
Original file line number	Diff line number	Diff line change
`@@ -23,7 +23,7 @@ namespace {`
`23`	`23`	`"Unable to create " #unary " layer from node: " << *n); \`
`24`	`24`	`\`
`25`	`25`	`unary->setName(util::node_info(n).c_str()); \`
`26`		`- auto out_tensor = associate_value_and_tensor(ctx, \`
	`26`	`+ auto out_tensor = ctx->AssociateValueAndTensor( \`
`27`	`27`	`n->outputs()[0], \`
`28`	`28`	`unary->getOutput(0)); \`
`29`	`29`	`LOG_DEBUG( \`