Implemented 1D conv, refactored pooling, 1D and special case for adap…

core/conversion/converters/BUILD

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

core/conversion/converters/converter_util.cpp

@@ -0,0 +1,53 @@
+#include "core/conversion/converters/converter_util.h"
+#include "core/conversion/converters/converters.h"
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+
+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) {
+      newDims = util::unsqueezeDims(newDims, trailing ? dim : 0, 1, use_zeros);
+    }
+
+    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->getOutput(0);
+  } else {
+    return tensor;
+  }
+}
+
+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) {
+      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->getOutput(0);
+  } else {
+    return tensor;
+  }
+}
+
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/conversion/converters/converter_util.h

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <map>
+#include <string>
+
+#include "core/conversion/conversionctx/ConversionCtx.h"
+#include "core/conversion/converters/Weights.h"
+#include "core/conversion/var/Var.h"
+#include "core/util/prelude.h"
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+
+  // 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
+} // namespace trtorch

core/conversion/converters/impl/batch_norm.cpp

@@ -1,3 +1,4 @@
+#include "core/conversion/converters/converter_util.h"
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
 #include "torch/torch.h"
@@ -40,17 +41,11 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       LOG_DEBUG("training disregarded");
       LOG_DEBUG("cudnn disregarded");
 
-      auto should_unpack = util::toVec(orig_shape).size() < 4;
-      if (should_unpack) {
-        // expand spatial dims from 1D to 2D
-        auto new_shape = util::toDimsTailPad(util::toVec(orig_shape), 4);
-        LOG_DEBUG(
-            "Input shape is less than 4D got: "
-            << orig_shape << ", inserting shuffle layer to reshape to 4D tensor shape: " << new_shape);
-        auto in_shuffle = ctx->net->addShuffle(*input);
-        in_shuffle->setReshapeDimensions(new_shape);
-        in_shuffle->setName(std::string("[Reshape input to " + util::toStr(new_shape) + ']').c_str());
-        input = in_shuffle->getOutput(0);
+      // Expand spatial dims from 1D to 2D if needed
+      bool expandDims = (orig_shape.nbDims < 4);
+
+      if (expandDims) {
+        input = addPadding(ctx, n, input, 4);
       }
 
       auto scale = gamma / torch::sqrt(var + eps);
@@ -63,16 +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 (should_unpack) {
-        LOG_DEBUG("Inserting shuffle layer to reshape to back to original shape: " << orig_shape);
-        auto out_shuffle = ctx->net->addShuffle(*out_tensor);
-        out_shuffle->setReshapeDimensions(orig_shape);
-        out_shuffle->setName(std::string("[Reshape output to " + util::toStr(orig_shape) + ']').c_str());
-        out_tensor = out_shuffle->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

@@ -1,7 +1,7 @@
-#include "torch/torch.h"
-
+#include "core/conversion/converters/converter_util.h"
 #include "core/conversion/converters/converters.h"
 #include "core/util/prelude.h"
+#include "torch/torch.h"
 
 namespace trtorch {
 namespace core {
@@ -14,15 +14,49 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
   auto in = args[0].ITensor(); // assumes non-static input Tensor
   auto w = Weights(ctx, args[1].unwrapToTensor());
   auto stride = util::toDims(args[3].unwrapToIntList());
-  LOG_DEBUG("stride: " << stride);
   auto padding = util::toDims(args[4].unwrapToIntList());
-  LOG_DEBUG("padding: " << padding);
   auto dilation = util::toDims(args[5].unwrapToIntList());
-  LOG_DEBUG("dilation: " << dilation);
   bool transposed = args[6].unwrapToBool();
   auto out_padding = util::toDims(args[7].unwrapToIntList());
-  LOG_DEBUG("out_padding: " << out_padding);
   int64_t groups = args[8].unwrapToInt();
+
+  auto dims = in->getDimensions();
+  auto orig_dims = dims;
+  LOG_DEBUG("Original input dims: " << orig_dims);
+
+  // Expand spatial dims from 1D to 2D if needed
+  bool expandDims = (orig_dims.nbDims < 4);
+  if (expandDims) {
+    in = addPadding(ctx, n, in, 4);
+    dims = in->getDimensions();
+  }
+  if (w.shape.nbDims < 4) {
+    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) {
+    stride = util::unsqueezeDims(stride, 1, 1);
+  }
+  if (dilation.nbDims==1) {
+    dilation = util::unsqueezeDims(dilation, 1, 1);
+  }
+  if (padding.nbDims==1) {
+    padding = util::unsqueezeDims(padding, 1, 0);
+  }
+  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);
+  LOG_DEBUG("padding: " << padding);
+  LOG_DEBUG("dilation: " << dilation);
+  LOG_DEBUG("out_padding: " << out_padding);
   LOG_DEBUG("groups: " << groups);
 
   nvinfer1::ILayer* new_layer;
@@ -31,12 +65,11 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
     if (args[2].IValue()->isTensor()) {
       bias = Weights(ctx, args[2].unwrapToTensor());
     } else {
-      bias = Weights(ctx, torch::zeros(args[1].unwrapToTensor().sizes()[1] * groups));
+      bias = Weights(ctx, torch::zeros(w.shape.d[1] * groups));
     }
 
     // shape of deconvolution's weight: [in, out/groups, ...]
-    auto deconv = ctx->net->addDeconvolutionNd(
-        *in, args[1].unwrapToTensor().sizes()[1] * groups, w.kernel_shape, w.data, bias.data);
+    auto deconv = ctx->net->addDeconvolutionNd(*in, w.shape.d[1] * groups, w.kernel_shape, w.data, bias.data);
     TRTORCH_CHECK(deconv, "Unable to create deconvolution layer from node: " << *n);
 
     deconv->setStrideNd(stride);
@@ -56,11 +89,11 @@ bool add_conv_deconv(ConversionCtx* ctx, const torch::jit::Node* n, args& args)
     if (args[2].IValue()->isTensor()) {
       bias = Weights(ctx, args[2].unwrapToTensor());
     } else {
-      bias = Weights(ctx, torch::zeros(args[1].unwrapToTensor().sizes()[0]));
+      bias = Weights(ctx, torch::zeros(w.shape.d[0]));
     }
 
     // shape of convolution's weight: [out, in/groups, ...]
-    auto conv = ctx->net->addConvolutionNd(*in, args[1].unwrapToTensor().sizes()[0], w.kernel_shape, w.data, bias.data);
+    auto conv = ctx->net->addConvolutionNd(*in, w.shape.d[0], w.kernel_shape, w.data, bias.data);
     TRTORCH_CHECK(conv, "Unable to create convolution layer from node: " << *n);
 
     conv->setStrideNd(stride);
@@ -71,9 +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());
-
-  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());