Remove redundent code in layer_norm and generalize layer_norm to accept gamma and beta as None.

Signed-off-by: Yu-Te Cheng <[email protected]>

Author: chengyute

core/conversion/converters/impl/layer_norm.cpp

@@ -81,23 +81,12 @@ auto layer_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
 
       /* Layer_Norm normalizes over last N dimensions.
          normalizaed_shape could be (C,H,W), (H,W), or (W). */
-
       auto normalized_shape = args[1].unwrapToIntList();
       auto normalized_shape_vec = util::toVec(util::toDims(normalized_shape));
 
-      torch::Tensor gamma, beta;
-      gamma = args[2].unwrapToTensor();
-      beta = args[3].unwrapToTensor();
-
-      // Remove batch dimension from input shape for expand_size, which will
-      // be used to create weights for addScaleNd later.
-      auto expand_size = shape;
-      expand_size.erase(expand_size.begin(), expand_size.begin() + 1);
-      auto gamma_expand = gamma.expand(expand_size);
-      auto beta_expand = beta.expand(expand_size);
-
       // Unwrap eps.
       auto eps = args[4].unwrapToDouble();
+
       LOG_DEBUG("cudnn disregarded");
 
       // Set up  axis_ask for E[x].
@@ -108,144 +97,89 @@ auto layer_norm_registrations TRTORCH_UNUSED = RegisterNodeConversionPatterns().
       LOG_DEBUG("Axis Mask for E[x]" << std::bitset<32>(axis_mask));
 
       // E[x]
-      auto mean_layer_expected = ctx->net->addReduce(*input, nvinfer1::ReduceOperation::kAVG, axis_mask, false);
-      TRTORCH_CHECK(mean_layer_expected, "Unable to create mean_layer_expected from node: " << *n);
-      mean_layer_expected->setName((util::node_info(n) + "_mean_expected").c_str());
-      auto mean_layer_expected_out = mean_layer_expected->getOutput(0);
-
-      // Expand output of E[x] to the same shape as original input.
-      c10::List<int64_t> repeats_expected;
-      for (size_t i = 0; i < shape.size(); i++) {
-        auto repeat = i > (shape.size() - normalized_shape_vec.size() - 1) ? shape[i] : 1;
-        repeats_expected.push_back(repeat);
-      }
-
-      int repeats_expected_rank = repeats_expected.size();
-      auto mean_layer_expected_out_dims = mean_layer_expected_out->getDimensions();
-      auto num_expand_dims_expected = repeats_expected_rank - mean_layer_expected_out_dims.nbDims;
-
-      if (num_expand_dims_expected > 0) {
-        nvinfer1::Dims reshape_expected_dims;
-        reshape_expected_dims.nbDims = repeats_expected.size();
-        for (int i = 0; i < num_expand_dims_expected; i++) {
-          reshape_expected_dims.d[repeats_expected.size() - 1 - i] = 1;
-        }
-        for (int i = 0; i < mean_layer_expected_out_dims.nbDims; i++) {
-          reshape_expected_dims.d[i] = mean_layer_expected_out_dims.d[i];
-        }
-        // Add a reshape layer to expand dims
-        auto reshape_layer_expected = ctx->net->addShuffle(*mean_layer_expected_out);
-        reshape_layer_expected->setReshapeDimensions(reshape_expected_dims);
-        mean_layer_expected_out = reshape_layer_expected->getOutput(0);
-      }
-
-      for (int i = repeats_expected.size() - 1; i >= 0; --i) {
-        std::vector<nvinfer1::ITensor*> tensors_vec;
-        for (int j = 0; j < repeats_expected[i]; j++) {
-          tensors_vec.push_back(mean_layer_expected_out);
-        }
-        auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
-        concat_layer->setAxis(i);
-        mean_layer_expected_out = concat_layer->getOutput(0);
-      }
+      auto mean_expected = ctx->net->addReduce(*input, nvinfer1::ReduceOperation::kAVG, axis_mask, true);
+      TRTORCH_CHECK(mean_expected, "Unable to create mean_expected from node: " << *n);
+      mean_expected->setName((util::node_info(n) + "_mean_expected").c_str());
+      auto mean_expected_out = mean_expected->getOutput(0);
 
       // X-E[x]
       auto sub = add_elementwise(
-          ctx,
-          nvinfer1::ElementWiseOperation::kSUB,
-          input,
-          mean_layer_expected_out,
-          (util::node_info(n) + "_sub").c_str());
-      TRTORCH_CHECK(sub, "Unable to create Add layer from node: " << *n);
+          ctx, nvinfer1::ElementWiseOperation::kSUB, input, mean_expected_out, (util::node_info(n) + "_sub").c_str());
+      TRTORCH_CHECK(sub, "Unable to create Sub layer from node: " << *n);
       sub->setName((util::node_info(n) + "_sub").c_str());
-      auto xsubmean = sub->getOutput(0);
+      auto xsubmean_out = sub->getOutput(0);
 
-      // Variance
+      // Variance = mean(pow(xsubmean,2))
       float pow_scalar = 2;
       auto exponent = tensor_to_const(ctx, torch::tensor({pow_scalar}));
       auto pow = add_elementwise(
-          ctx, nvinfer1::ElementWiseOperation::kPOW, xsubmean, exponent, (util::node_info(n) + "_pow").c_str());
-      TRTORCH_CHECK(pow, "Unable to create Power layer from node: " << *n);
+          ctx, nvinfer1::ElementWiseOperation::kPOW, xsubmean_out, exponent, (util::node_info(n) + "_pow").c_str());
+      TRTORCH_CHECK(pow, "Unable to create Pow layer from node: " << *n);
       pow->setName((util::node_info(n) + "_pow").c_str());
       auto pow_out = pow->getOutput(0);
 
-      auto mean_layer_var = ctx->net->addReduce(*pow_out, nvinfer1::ReduceOperation::kAVG, axis_mask, false);
-      TRTORCH_CHECK(mean_layer_var, "Unable to create mean_layer_var from node: " << *n);
-      mean_layer_var->setName((util::node_info(n) + "_mean_var").c_str());
-      auto mean_layer_var_out = mean_layer_var->getOutput(0);
-
-      // Expand output of mean_layer_var to the same shape as original
-      // input.
-      c10::List<int64_t> repeats_var;
-      for (size_t i = 0; i < shape.size(); i++) {
-        auto repeat = i > (shape.size() - normalized_shape_vec.size() - 1) ? shape[i] : 1;
-        repeats_var.push_back(repeat);
-      }
-
-      int repeats_var_rank = repeats_var.size();
-      auto mean_layer_var_out_dims = mean_layer_var_out->getDimensions();
-      auto num_expand_dims_var = repeats_var_rank - mean_layer_var_out_dims.nbDims;
-
-      if (num_expand_dims_var > 0) {
-        nvinfer1::Dims reshape_dims_var;
-        reshape_dims_var.nbDims = repeats_var.size();
-        for (int i = 0; i < num_expand_dims_var; i++) {
-          reshape_dims_var.d[repeats_var.size() - 1 - i] = 1;
-        }
-        for (int i = 0; i < mean_layer_var_out_dims.nbDims; i++) {
-          reshape_dims_var.d[i] = mean_layer_var_out_dims.d[i];
-        }
-
-        // Add a reshape layer to expand dims
-        auto reshape_layer_var = ctx->net->addShuffle(*mean_layer_var_out);
-        reshape_layer_var->setReshapeDimensions(reshape_dims_var);
-        mean_layer_var_out = reshape_layer_var->getOutput(0);
-      }
+      auto mean_var = ctx->net->addReduce(*pow_out, nvinfer1::ReduceOperation::kAVG, axis_mask, true);
+      TRTORCH_CHECK(mean_var, "Unable to create mean_var from node: " << *n);
+      mean_var->setName((util::node_info(n) + "_mean_var").c_str());
+      auto mean_var_out = mean_var->getOutput(0);
 
-      for (int i = repeats_var.size() - 1; i >= 0; --i) {
-        std::vector<nvinfer1::ITensor*> tensors_vec;
-        for (int j = 0; j < repeats_var[i]; j++) {
-          tensors_vec.push_back(mean_layer_var_out);
-        }
-        auto concat_layer = ctx->net->addConcatenation(tensors_vec.data(), tensors_vec.size());
-        concat_layer->setAxis(i);
-        mean_layer_var_out = concat_layer->getOutput(0);
-      }
-
-      // add eps
+      // Variance + eps
       auto eps_tensor = tensor_to_const(ctx, torch::tensor({eps}));
       auto add = add_elementwise(
-          ctx,
-          nvinfer1::ElementWiseOperation::kSUM,
-          mean_layer_var_out,
-          eps_tensor,
-          (util::node_info(n) + "_add").c_str());
+          ctx, nvinfer1::ElementWiseOperation::kSUM, mean_var_out, eps_tensor, (util::node_info(n) + "_add").c_str());
       TRTORCH_CHECK(add, "Unable to create Add layer from node: " << *n);
       add->setName((util::node_info(n) + "_add").c_str());
       auto add_out = add->getOutput(0);
 
-      // add Unary layer for sqrt((var + eps))
-      auto unary = ctx->net->addUnary(*add_out, nvinfer1::UnaryOperation::kSQRT);
-      TRTORCH_CHECK(unary, "Unable to create unary layer from node: " << *n);
-      unary->setName((util::node_info(n) + "_unary_sqrt").c_str());
-      auto unary_out = unary->getOutput(0);
+      // SQRT((Var + eps))
+      auto sqrt = ctx->net->addUnary(*add_out, nvinfer1::UnaryOperation::kSQRT);
+      TRTORCH_CHECK(sqrt, "Unable to create unary(sqrt) from node: " << *n);
+      sqrt->setName((util::node_info(n) + "_sqrt").c_str());
+      auto sqrt_out = sqrt->getOutput(0);
 
       // (x - E[x]) / sqrt((var + eps))
       auto div = add_elementwise(
-          ctx, nvinfer1::ElementWiseOperation::kDIV, xsubmean, unary_out, (util::node_info(n) + "_div").c_str());
+          ctx, nvinfer1::ElementWiseOperation::kDIV, xsubmean_out, sqrt_out, (util::node_info(n) + "_div").c_str());
       TRTORCH_CHECK(div, "Unable to create div layer from node: " << *n);
       div->setName((util::node_info(n) + "_div").c_str());
       auto div_out = div->getOutput(0);
 
+      if (!args[2].IValue()->isTensor() && !args[3].IValue()->isTensor()) {
+        ctx->AssociateValueAndTensor(n->outputs()[0], div_out);
+        return true;
+      }
+
+      // Remove batch dimension from input shape for expand_size, which will
+      // be used to create weights for addScaleNd later.
+      auto expand_size = shape;
+      expand_size.erase(expand_size.begin(), expand_size.begin() + 1);
+
       // Set up gamma_weights and beta_weights from gamma_expand and
-      // beta_expand
-      auto gamma_weights = Weights(ctx, gamma_expand);
-      auto beta_weights = Weights(ctx, beta_expand);
+      // beta_expand.
+      auto gamma_weights = Weights(ctx, at::ones(expand_size));
+      auto beta_weights = Weights(ctx, at::zeros(expand_size));
+
+      if (args[2].IValue()->isTensor()) {
+        torch::Tensor gamma;
+        gamma = args[2].unwrapToTensor();
+        auto gamma_expand = gamma.expand(expand_size);
+        gamma_weights = Weights(ctx, gamma_expand);
+      } else {
+        gamma_weights = Weights(ctx, at::ones(expand_size));
+      }
 
-      auto power = Weights(ctx, at::ones_like(gamma_expand));
+      if (args[3].IValue()->isTensor()) {
+        torch::Tensor beta;
+        beta = args[3].unwrapToTensor();
+        auto beta_expand = beta.expand(expand_size);
+        beta_weights = Weights(ctx, beta_expand);
+      } else {
+        beta_weights = Weights(ctx, at::zeros(expand_size));
+      }
+
+      auto power = Weights(ctx, at::ones(expand_size));
       auto scale_nd = ctx->net->addScaleNd(
           *div_out, nvinfer1::ScaleMode::kELEMENTWISE, beta_weights.data, gamma_weights.data, power.data, 1);
-
       scale_nd->setName((util::node_info(n) + "_scale_nd").c_str());
       auto scale_nd_out = scale_nd->getOutput(0);
 

tests/core/conversion/converters/BUILD

@@ -36,7 +36,7 @@ converter_test(
 )
 
 converter_test(
-  name = "test_layer_norm",
+    name = "test_layer_norm",
 )
 
 converter_test(
@@ -114,6 +114,7 @@ test_suite(
         ":test_element_wise",
         ":test_expand",
         ":test_interpolate",
+        ":test_layer_norm",
         ":test_linear",
         ":test_lstm_cell",
         ":test_matrix_multiply",

tests/core/conversion/converters/test_layer_norm.cpp

@@ -4,6 +4,36 @@
 #include "tests/util/util.h"
 #include "torch/csrc/jit/ir/irparser.h"
 
+TEST(Converters, ATenLayerNormConvertsCorrectlyLast3DimsNoGammaBeta) {
+  const auto graph = R"IR(
+      graph(%0 : Tensor):
+        %gamma : None = prim::Constant()
+        %beta : None = prim::Constant()
+        %1: int = prim::Constant[value=3]()
+        %2: int = prim::Constant[value=100]()
+        %3: int = prim::Constant[value=100]()
+        %4 : int[] = prim::ListConstruct(%1, %2, %3)
+        %7 : bool = prim::Constant[value=0]()
+        %8 : float = prim::Constant[value=1.0000000000000001e-05]()
+        %9 : Tensor = aten::layer_norm(%0, %4, %gamma, %beta, %8, %7)
+        return (%9))IR";
+
+  auto g = std::make_shared<torch::jit::Graph>();
+  torch::jit::parseIR(graph, g.get());
+
+  auto in = at::randint(1, 10, {4, 3, 100, 100}, {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));
+ 
+}
+
+
 TEST(Converters, ATenLayerNormConvertsCorrectlyLast3Dims) {
   const auto graph = R"IR(
       graph(%0 : Tensor,