Merge pull request #405 from guoruoqian/batchnorm_fix_bug

Fix bug when batchnorm should unpack

Author: narendasan

core/conversion/converters/impl/batch_norm.cpp

@@ -43,7 +43,7 @@ auto batch_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       auto should_unpack = util::toVec(orig_shape).size() < 4;
       if (should_unpack) {
         // expand spatial dims from 1D to 2D
-        auto new_shape = util::toDimsPad(util::toVec(orig_shape), 4);
+        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);

core/util/trt_util.cpp

@@ -86,6 +86,30 @@ nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to) {
   return dims;
 }
 
+nvinfer1::Dims toDimsTailPad(c10::IntArrayRef l, uint64_t pad_to) {
+  if (l.size() > pad_to) {
+    LOG_DEBUG(
+        "Requested padding of dimensions to " << pad_to << " but found " << l.size()
+                                              << " dimensions, not going to pad");
+    return toDims(l);
+  }
+
+  TRTORCH_CHECK(
+      pad_to <= nvinfer1::Dims::MAX_DIMS,
+      "The list requested to be converted to nvinfer1::Dims exceeds the max number of dimensions for TensorRT");
+
+  nvinfer1::Dims dims;
+  dims.nbDims = pad_to;
+  for (size_t i = 0; i < l.size(); i++) {
+    dims.d[i] = l[i];
+  }
+
+  for (size_t i = pad_to - l.size(); i < pad_to; i++) {
+    dims.d[i] = 1;
+  }
+  return dims;
+}
+
 nvinfer1::Dims toDims(c10::IntArrayRef l) {
   TRTORCH_CHECK(
       l.size() <= nvinfer1::Dims::MAX_DIMS,
@@ -136,6 +160,30 @@ nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to) {
   return dims;
 }
 
+nvinfer1::Dims toDimsTailPad(c10::List<int64_t> l, uint64_t pad_to) {
+  if (l.size() > pad_to) {
+    LOG_DEBUG(
+        "Requested padding of dimensions to " << pad_to << " but found " << l.size()
+                                              << " dimensions, not going to pad");
+    return toDims(l);
+  }
+
+  TRTORCH_CHECK(
+      pad_to <= nvinfer1::Dims::MAX_DIMS,
+      "The list requested to be converted to nvinfer1::Dims exceeds the max number of dimensions for TensorRT");
+
+  nvinfer1::Dims dims;
+  dims.nbDims = pad_to;
+  for (size_t i = 0; i < l.size(); i++) {
+    dims.d[i] = l[i];
+  }
+
+  for (size_t i = pad_to - l.size(); i < pad_to; i++) {
+    dims.d[i] = 1;
+  }
+  return dims;
+}
+
 nvinfer1::Dims unpadDims(const nvinfer1::Dims& d) {
   nvinfer1::Dims dims;
 
@@ -304,4 +352,4 @@ c10::optional<nvinfer1::DataType> toTRTDataType(caffe2::TypeMeta dtype) {
 
 } // namespace util
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch

core/util/trt_util.h

@@ -92,6 +92,8 @@ int64_t volume(const nvinfer1::Dims& d);
 bool broadcastable(nvinfer1::Dims a, nvinfer1::Dims b, bool multidirectional = true);
 nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to);
 nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to);
+nvinfer1::Dims toDimsTailPad(c10::IntArrayRef l, uint64_t pad_to);
+nvinfer1::Dims toDimsTailPad(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);
@@ -110,4 +112,4 @@ const std::unordered_map<at::ScalarType, nvinfer1::DataType>& get_aten_trt_type_
 
 } // namespace util
 } // namespace core
-} // namespace trtorch
+} // namespace trtorch

tests/core/conversion/converters/test_batch_norm.cpp

@@ -34,3 +34,34 @@ TEST(Converters, ATenBatchNormConvertsCorrectly) {
 
   ASSERT_TRUE(trtorch::tests::util::almostEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), 2e-6));
 }
+
+TEST(Converters, ATenBatchNormShouldUnpackConvertsCorrectly) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor,
+            %1: Float(5:1),
+            %2: Float(5:1),
+            %3: Float(5:1),
+            %4: Float(5:1)):
+        %5 : bool = prim::Constant[value=0]()
+        %6 : float = prim::Constant[value=1.0000000000000001e-05]()
+        %7 : float = prim::Constant[value=0.10000000000000001]()
+        %8 : Tensor = aten::batch_norm(%0, %1, %2, %3, %4, %5, %6, %7, %5)
+        return (%8))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, &*g);
+
+  auto in = at::randint(1, 10, {3, 5}, {at::kCUDA});
+  auto gamma = at::randint(1, 10, {5}, {at::kCUDA});
+  auto beta = at::randint(1, 10, {5}, {at::kCUDA});
+  auto mean = at::randint(1, 10, {5}, {at::kCUDA});
+  auto var = at::randint(1, 10, {5}, {at::kCUDA});
+
+  auto params = trtorch::core::conversion::get_named_params(g->inputs(), {gamma, beta, mean, var});
+  auto jit_results = trtorch::tests::util::RunGraph(g, params, {in});
+
+  params = trtorch::core::conversion::get_named_params(g->inputs(), {gamma, beta, mean, var});
+  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));
+}