Addressing code review comments

Signed-off-by: Boris Fomitchev <[email protected]>

Author: borisfom

core/conversion/converters/BUILD

@@ -30,11 +30,11 @@ cc_library(
     name = "converters",
     hdrs = [
         "converters.h",
-	"converter_util.h",
+        "converter_util.h",
     ],
     srcs = [
         "NodeConverterRegistry.cpp",
-	"converter_util.cpp",	
+        "converter_util.cpp",   
         "impl/activation.cpp",
         "impl/batch_norm.cpp",
         "impl/concat.cpp",

core/conversion/converters/converter_util.cpp

@@ -7,43 +7,44 @@ namespace core {
 namespace conversion {
 namespace converters {
 
-nvinfer1::ILayer* addPaddingLayer(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing, bool use_zeros) {
+nvinfer1::ITensor* addPadding(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing, bool use_zeros) {
   const auto dims = tensor->getDimensions();
 
   if (dims.nbDims < nDim) {
     auto newDims = dims;
-    for (int dim = dims.nbDims; dim < nDim; ++dim)
+    for (int dim = dims.nbDims; dim < nDim; ++dim) {
       newDims = util::unsqueezeDims(newDims, trailing ? dim : 0, 1, use_zeros);
-    LOG_DEBUG(
-        "Input shape is less than 4D got: " << dims
-                                            << ", inserting shuffle layer to reshape to 4D tensor shape: " << newDims);
+    }
 
     LOG_DEBUG("Original shape: " << dims << ", reshaping to: " << newDims);
     auto shuffle_layer = ctx->net->addShuffle(*tensor);
     TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer");
     shuffle_layer->setReshapeDimensions(newDims);
     shuffle_layer->setZeroIsPlaceholder(use_zeros);
     shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims) + ']').c_str());
-    return shuffle_layer;
-  } else
+    return shuffle_layer->getOutput(0);
+  } else {
     return nullptr;
+  }
 }
 
-nvinfer1::ILayer* addUnpaddingLayer(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing, bool use_zeros) {
+nvinfer1::ITensor* addUnpadding(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing, bool use_zeros) {
   const auto dims = tensor->getDimensions();
   if (dims.nbDims > nDim) {
     auto newDims = dims;
-    for (int dim = dims.nbDims; dim > nDim; --dim)
+    for (int dim = dims.nbDims; dim > nDim; --dim) {
       newDims = util::squeezeDims(newDims, trailing ? dim - 1 : 0);
+    }
     LOG_DEBUG("Original shape: " << dims << ", reshaping to: " << newDims);
     auto shuffle_layer = ctx->net->addShuffle(*tensor);
     TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer");
     shuffle_layer->setReshapeDimensions(newDims);
     shuffle_layer->setZeroIsPlaceholder(use_zeros);
     shuffle_layer->setName((util::node_info(n) + " [Reshape to " + util::toStr(newDims)).c_str() + ']');
-    return shuffle_layer;
-  } else
-    return nullptr;
+    return shuffle_layer->getOutput(0);
+  } else {
+    return tensor;
+  }
 }
 
 } // namespace converters

core/conversion/converters/converter_util.h

@@ -12,8 +12,17 @@ namespace trtorch {
 namespace core {
 namespace conversion {
 namespace converters {
-nvinfer1::ILayer* addPaddingLayer(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing=true, bool use_zeros=true);
-nvinfer1::ILayer* addUnpaddingLayer(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing=true, bool use_zeros=true);
+  
+  // If nDim < tensor size, adds shuffle layer to pad tensor with 1s (at the end if trailing) and returns (nDim-dimensional) shuffle layer's output.
+  // Otherwise, does nothing and passes tensor through. 
+  // use _zeros controls whether we should be using 0 instead of -1 on the shape.
+  nvinfer1::ITensor* addPadding(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing=true, bool use_zeros=true);
+  
+  // If nDim < tensor size, adds shuffle layer to un-pad tensor (at the end if trailing) and returns (nDim-dimensional) shuffle layer's output
+  // Otherwise, does nothing and passes tensor through. 
+  // use _zeros controls whether we should be using 0 instead of -1 on the shape.
+  nvinfer1::ITensor* addUnpadding(ConversionCtx* ctx, const torch::jit::Node* n, nvinfer1::ITensor* tensor, int nDim, bool trailing=true, bool use_zeros=true);
+  
 } // namespace converters
 } // namespace conversion
 } // namespace core

core/conversion/converters/impl/batch_norm.cpp

@@ -41,9 +41,11 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       LOG_DEBUG("training disregarded");
       LOG_DEBUG("cudnn disregarded");
 
-      auto expandDims = addPaddingLayer(ctx, n, input, 4);
+      // Expand spatial dims from 1D to 2D if needed
+      bool expandDims = (orig_shape.nbDims < 4);
+      
       if (expandDims) {
-        input = expandDims->getOutput(0);
+        input = addPadding(ctx, n, input, 4);
       }
 
       auto scale = gamma / torch::sqrt(var + eps);
@@ -56,13 +58,8 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       auto bn = ctx->net->addScaleNd(
           *input, nvinfer1::ScaleMode::kCHANNEL, bias_weights.data, scale_weights.data, power.data, 1);
       bn->setName(util::node_info(n).c_str());
-      auto out_tensor = bn->getOutput(0);
-
-      if (expandDims) {
-        LOG_DEBUG("Inserting shuffle layer to reshape to back to original shape: " << orig_shape);
-        auto new_layer = addUnpaddingLayer(ctx, n, out_tensor, orig_shape.nbDims);
-        out_tensor = new_layer->getOutput(0);
-      }
+      // Un-pad bn output if needed
+      auto out_tensor = addUnpadding(ctx, n, bn->getOutput(0), orig_shape.nbDims);
       ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
       return true;
     }});

core/conversion/converters/impl/conv_deconv.cpp

@@ -25,29 +25,32 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
   LOG_DEBUG("Original input dims: " << orig_dims);
 
   // Expand spatial dims from 1D to 2D if needed
-  auto expandDims = addPaddingLayer(ctx, n, in, 4);
+  bool expandDims = (orig_dims.nbDims < 4);
   if (expandDims) {
-    auto tensorPtr = expandDims->getOutput(0);
-    assert(tensorPtr);
-    dims = tensorPtr->getDimensions();
-    in = tensorPtr;
+    in = addPadding(ctx, n, in, 4);
+    dims = in->getDimensions();
   }
   if (w.shape.nbDims < 4) {
-    for (int i = w.shape.nbDims; i < 4; ++i)
+    for (int i = w.shape.nbDims; i < 4; ++i) {
       w.shape.d[i] = 1;
+    }
     w.shape.nbDims = 4;
     w.kernel_shape.nbDims = 2;
     w.kernel_shape.d[1] = 1;
   }
-  if (stride.nbDims==1)
+  if (stride.nbDims==1) {
     stride = util::unsqueezeDims(stride, 1, 1);
-  if (dilation.nbDims==1)
+  }
+  if (dilation.nbDims==1) {
     dilation = util::unsqueezeDims(dilation, 1, 1);
-  if (padding.nbDims==1)  
+  }
+  if (padding.nbDims==1) {
     padding = util::unsqueezeDims(padding, 1, 0);
-  if (out_padding.nbDims==1)
+  }
+  if (out_padding.nbDims==1) {
     out_padding = util::unsqueezeDims(out_padding, 1, 0);
-
+  }
+  
   LOG_DEBUG("Input dims: " << dims);
   LOG_DEBUG("Weights: " << w);
   LOG_DEBUG("stride: " << stride);
@@ -101,14 +104,13 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
     conv->setNbGroups(groups);
     new_layer = conv;
   }
+  
   new_layer->setName(util::node_info(n).c_str());
-
-  if (expandDims) {
-    // Un-expand spatial dims back to 1D
-    new_layer = addUnpaddingLayer(ctx, n, new_layer->getOutput(0), orig_dims.nbDims);
-  }
-
-  auto out = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+  
+  // Un-expand spatial dims back to 1D if needed
+  auto out = addUnpadding(ctx, n, new_layer->getOutput(0), orig_dims.nbDims);
+  
+  ctx->AssociateValueAndTensor(n->outputs()[0], out);
 
   LOG_DEBUG("Output tensor shape: " << out->getDimensions());
 

core/conversion/converters/impl/pooling.cpp

@@ -29,22 +29,23 @@ bool GlobalPoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args&
 bool AdaptivePoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,  nvinfer1::PoolingType pool_type) {
   auto in = args[0].ITensorOrFreeze(ctx);
   auto out_size = util::toDims(args[1].unwrapToIntList());
-  bool shuffle_back = false;
 
   // Corner case: when out dimension is all ones, replace with simpler operation
   if (out_size.d[0] == 1 && (out_size.nbDims < 2 || out_size.d[1] == 1 ) && (out_size.nbDims < 3 || out_size.d[2] == 1 ))  {
     return GlobalPoolingConverter(ctx, n, args, pool_type);
   }
 
-  auto shuffle = addPaddingLayer(ctx, n, in, 4, false, false);
-  if (shuffle) {
-    in = shuffle->getOutput(0);
+  auto orig_dims = in->getDimensions();
+  bool expandDims = (orig_dims.nbDims < 4); 
+ 
+  if (expandDims) {
+    in = addPadding(ctx, n, in, 4, false, false);
   }
 
   if (out_size.nbDims == 1) {
     out_size = util::unsqueezeDims(out_size, 0, 1);
-    shuffle_back = true;
   }
+  
   auto in_shape = util::toVec(in->getDimensions());
   nvinfer1::ILayer* new_layer = nullptr;
 
@@ -91,13 +92,10 @@ bool AdaptivePoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, arg
     new_layer = pooling_layer;
   }
 
-  new_layer->setName(util::node_info(n).c_str());
-
-  if (shuffle_back ) {
-    new_layer = addUnpaddingLayer(ctx, n, new_layer->getOutput(0), 3, false, false);
-  }
+  new_layer->setName(util::node_info(n).c_str()); 
+  auto layer_output = addUnpadding(ctx, n, new_layer->getOutput(0), orig_dims.nbDims, false, false);
 
-  auto layer_output = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));
+  ctx->AssociateValueAndTensor(n->outputs()[0], layer_output);
   LOG_DEBUG("Output tensor shape: " << layer_output->getDimensions());
 
   return true;
@@ -107,10 +105,11 @@ bool PoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,
   auto in = args[0].ITensorOrFreeze(ctx);
 
   // Max Pool needs at least 4D input
-  auto shuffle = addPaddingLayer(ctx, n, in, 4, false, true);
+  auto orig_dims = in->getDimensions();
+  bool expandDims = (orig_dims.nbDims < 4);
 
-  if (shuffle) {
-    in = shuffle->getOutput(0);
+  if (expandDims) {
+    in = addPadding(ctx, n, in, 4, false, true);
   }
 
   auto kernel_size = util::toDims(args[1].unwrapToIntList());
@@ -121,19 +120,18 @@ bool PoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,
     stride = util::toDims(args[1].unwrapToIntList());
   }
 
-  bool shuffle_back = false;
   if (kernel_size.nbDims == 1) {
     kernel_size = util::unsqueezeDims(kernel_size, 0, 1);
-    if (shuffle)
-      shuffle_back = true;
     LOG_DEBUG("kernel_size.nbDims < 2, padding:" << kernel_size);
     LOG_DEBUG("kernel_size: " << kernel_size);
   }
-  if (padding.nbDims == 1)
-    padding = util::unsqueezeDims(padding, 0, 0); 
-  if (stride.nbDims == 1)
+  if (padding.nbDims == 1) {
+    padding = util::unsqueezeDims(padding, 0, 0);
+  }
+  if (stride.nbDims == 1) {
     stride = util::unsqueezeDims(stride, 0, 1);
-
+  }
+  
   LOG_DEBUG("kernel_size: " << kernel_size);
   LOG_DEBUG("padding: " << padding);
   LOG_DEBUG("stride: " << stride);
@@ -161,37 +159,28 @@ bool PoolingConverter(ConversionCtx* ctx, const torch::jit::Node* n, args& args,
     new_layer = ctx->net->addPoolingNd(*in, pool_type, kernel_size);
     TRTORCH_CHECK(new_layer, "Unable to create Avg Pooling layer from node: " << *n);
     new_layer->setAverageCountExcludesPadding(!count_inlcude_pad);
-    // if (!(args[6].IValue()->isNone())) {
-    //  LOG_WARNING("Divisor override is now handled by Avg Pooling Converter");
-    // }
   } else {
     TRTORCH_ASSERT(0, "Unsupported pool mode!");
   }
 
+  auto padding_mode =
+      ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
+  
   new_layer->setName(util::node_info(n).c_str());
+  new_layer->setPaddingMode(padding_mode);
   new_layer->setPaddingNd(padding);
+  new_layer->setStrideNd(stride);
+  
   if (stride.nbDims != 2 && ctx->settings.device.device_type == nvinfer1::DeviceType::kDLA) {
     if (!ctx->settings.device.allow_gpu_fallback) {
       TRTORCH_THROW_ERROR("DLA Pooling stride is limited to 2D, allow GPU fallback");
     } else {
       LOG_WARNING("DLA Pooling stride is limited to 2D, will run on GPU");
     }
   }
-  new_layer->setStrideNd(stride);
-
-  auto padding_mode =
-      ceil_mode ? nvinfer1::PaddingMode::kEXPLICIT_ROUND_UP : nvinfer1::PaddingMode::kEXPLICIT_ROUND_DOWN;
-  new_layer->setPaddingMode(padding_mode);
-
-  new_layer->setName(util::node_info(n).c_str());
-
-  nvinfer1::ILayer* out_layer = new_layer;
 
-  if (shuffle_back) {
-    out_layer = addUnpaddingLayer(ctx, n, new_layer->getOutput(0), 3, false, true);
-  }
-
-  auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], out_layer->getOutput(0));
+  auto out_tensor = addUnpadding(ctx, n, new_layer->getOutput(0), orig_dims.nbDims, false, true);
+  ctx->AssociateValueAndTensor(n->outputs()[0], out_tensor);
 
   LOG_DEBUG("Output tensor shape: " << out_tensor->getDimensions());
   return true;

core/conversion/evaluators/prim.cpp

@@ -48,13 +48,7 @@ auto prim_registrations =
                           c10::List<int64_t> list;
                           list.reserve(num_inputs);
                           for (auto in : n->inputs()) {
-                            if (args.at(in).IValue()->isInt()) {
-                              list.emplace_back(std::move(args.at(in).unwrapToInt()));
-                            } else if (args.at(in).IValue()->isTuple()) {
-                              auto unpack_tuple = args.at(in).IValue()->toTuple();
-                              for (size_t j = 0; j < unpack_tuple->elements().size(); ++j)
-                                list.emplace_back(unpack_tuple->elements()[j].toInt());
-                            }
+                            list.emplace_back(std::move(args.at(in).unwrapToInt()));
                           }
                           return c10::optional<torch::jit::IValue>(std::move(torch::jit::IValue(list)));
                         } else if (torch::jit::FloatType::get() == lt->getElementType()) {

tests/core/conversion/converters/test_conv_deconv.cpp

@@ -98,8 +98,8 @@ TEST(Converters, ATenConvolution1dConvertsCorrectly) {
   auto g = std::make_shared<torch::jit::Graph>();
   torch::jit::parseIR(graph, g.get());
 
-  auto in = at::randint(1, 2, {1, 1, 3, 3}, {at::kCUDA});
-  auto w = at::randint(1, 2, {4, 1, 2, 2}, {at::kCUDA});
+  auto in = at::randint(1, 2, {1, 3, 3}, {at::kCUDA});
+  auto w = at::randint(1, 2, {4, 3, 3}, {at::kCUDA});
 
   auto jit_in = at::clone(in);
   auto jit_w = at::clone(w);