Add support for aten::squeeze, aten::unsqueeze

core/conversion/converters/BUILD

@@ -48,8 +48,10 @@ cc_library(
         "impl/unary.cpp",
         "impl/interpolate.cpp",
         "impl/select.cpp",
+        "impl/squeeze.cpp",
         "impl/stack.cpp",
-        "impl/lstm_cell.cpp"
+        "impl/lstm_cell.cpp",
+        "impl/unsqueeze.cpp",
     ],
     deps = [
         "@tensorrt//:nvinfer",

core/conversion/converters/impl/squeeze.cpp

@@ -0,0 +1,44 @@
+#include "NvInfer.h"
+#include "core/conversion/converters/converters.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/prelude.h"
+#include "torch/torch.h"
+
+#include <ATen/ATen.h>
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto squeeze_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::squeeze.dim(Tensor(a) self, int dim) -> (Tensor(a))",
+     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       auto self = args[0].ITensorOrFreeze(ctx);
+       auto dim = args[1].unwrapToInt();
+
+       auto selfDim = util::toVec(self->getDimensions());
+       if (dim < 0) {
+         dim = selfDim.size() + dim;
+       }
+
+       auto shuffle_layer = ctx->net->addShuffle(*self);
+       TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+       shuffle_layer->setReshapeDimensions(util::squeezeDims(self->getDimensions(), dim));
+
+       auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_layer->getOutput(0));
+
+       LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+       return true;
+     }});
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/conversion/converters/impl/unsqueeze.cpp

@@ -0,0 +1,44 @@
+#include "NvInfer.h"
+#include "core/conversion/converters/converters.h"
+#include "core/conversion/tensorcontainer/TensorContainer.h"
+#include "core/util/prelude.h"
+#include "torch/torch.h"
+
+#include <ATen/ATen.h>
+#include <vector>
+
+namespace trtorch {
+namespace core {
+namespace conversion {
+namespace converters {
+namespace impl {
+namespace {
+
+auto unsqueeze_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().pattern(
+    {"aten::unsqueeze(Tensor(a) self, int dim) -> (Tensor(a))",
+     [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
+       auto self = args[0].ITensorOrFreeze(ctx);
+       auto dim = args[1].unwrapToInt();
+
+       auto selfDim = util::toVec(self->getDimensions());
+       if (dim < 0) {
+         dim = selfDim.size() + dim;
+       }
+
+       auto shuffle_layer = ctx->net->addShuffle(*self);
+       TRTORCH_CHECK(shuffle_layer, "Unable to create shuffle layer from node: " << *n);
+       shuffle_layer->setReshapeDimensions(util::unsqueezeDims(self->getDimensions(), dim));
+
+       auto out = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle_layer->getOutput(0));
+
+       LOG_DEBUG("Output tensor shape: " << out->getDimensions());
+
+       return true;
+     }});
+
+} // namespace
+} // namespace impl
+} // namespace converters
+} // namespace conversion
+} // namespace core
+} // namespace trtorch

core/util/trt_util.cpp

@@ -186,6 +186,34 @@ nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos) {
   return dims;
 }
 
+nvinfer1::Dims squeezeDims(const nvinfer1::Dims& d, int pos) {
+  // acceptable range for pos is [0, d.nbDims]
+  TRTORCH_ASSERT(pos >= 0 && pos <= d.nbDims, "ERROR: Index to squeeze is out of bounds.");
+
+  nvinfer1::Dims dims;
+
+  int i = 0;
+  int j = 0;
+
+  while (i <= d.nbDims) {
+    if (j != pos) {
+      dims.d[j] = d.d[i];
+    } else {
+      // add new dimension at pos
+      i++;
+      if (i <= d.nbDims) {
+        dims.d[j] = d.d[i];
+      }
+    }
+    i++;
+    j++;
+  }
+
+  dims.nbDims = d.nbDims - 1;
+
+  return dims;
+}
+
 std::vector<int64_t> toVec(nvinfer1::Dims d) {
   std::vector<int64_t> dims;
   for (int i = 0; i < d.nbDims; i++) {

core/util/trt_util.h

@@ -94,6 +94,7 @@ nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to);
 nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to);
 nvinfer1::Dims unpadDims(const nvinfer1::Dims& d);
 nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos);
+nvinfer1::Dims squeezeDims(const nvinfer1::Dims& d, int pos);
 nvinfer1::Dims toDims(c10::IntArrayRef l);
 nvinfer1::Dims toDims(c10::List<int64_t> l);
 nvinfer1::DimsHW toDimsHW(c10::List<int64_t> l);

tests/core/conversion/converters/BUILD

@@ -71,6 +71,14 @@ converter_test(
   name = "test_lstm_cell"
 )
 
+converter_test(
+  name = "test_unsqueeze"
+)
+
+converter_test(
+  name = "test_squeeze"
+)
+
 test_suite(
   name = "converter_tests",
   tests = [
@@ -88,6 +96,8 @@ test_suite(
     ":test_interpolate",
     ":test_select",
     ":test_stack",
-    ":test_lstm_cell"
+    ":test_lstm_cell",
+    ":test_unsqueeze",
+    ":test_squeeze"
   ]
 )

tests/core/conversion/converters/test_squeeze.cpp

@@ -0,0 +1,26 @@
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Converters, ATenSqueezeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : int = prim::Constant[value=1]()
+        %2 : Tensor = aten::squeeze(%0, %1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {2, 1, 3, 3}, {at::kCUDA});
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}

tests/core/conversion/converters/test_unsqueeze.cpp

@@ -0,0 +1,26 @@
+#include <string>
+#include "core/compiler.h"
+#include "gtest/gtest.h"
+#include "tests/util/util.h"
+#include "torch/csrc/jit/ir/irparser.h"
+
+TEST(Converters, ATenUnsqueezeConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %1 : int = prim::Constant[value=2]()
+        %2 : Tensor = aten::unsqueeze(%0, %1)
+        return (%2))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {2, 3, 3}, {at::kCUDA});
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {});
+  auto trt_results = trtorch::tests::util::RunGraphEngine(g, params, {in});
+
+  ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
+}