support dynamic shape for batch_norm and batch_norm_ op (#68353)

zeroRains · web-flow · commit bb75a93f167c · 2024-09-23T09:37:29.000+08:00
diff --git a/paddle/fluid/primitive/base/decomp_trans.cc b/paddle/fluid/primitive/base/decomp_trans.cc
@@ -49,12 +49,8 @@ std::unordered_set<std::string> decomp_op_contain_none = {
     "pd_op.instance_norm",
 };
 //
-std::unordered_set<std::string> dynamic_shape_blacklist = {"pd_op.squeeze",
-                                                           "pd_op.unsqueeze",
-                                                           "pd_op.batch_norm",
-                                                           "pd_op.batch_norm_",
-                                                           "pd_op.bmm",
-                                                           "pd_op.flatten"};
+std::unordered_set<std::string> dynamic_shape_blacklist = {
+    "pd_op.squeeze", "pd_op.unsqueeze", "pd_op.bmm", "pd_op.flatten"};
 
 namespace {
 std::set<std::string> StringSplit(const std::string& str) {
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -371,9 +371,7 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> batch_norm_decomp(
   if (need_cast) {
     x_cast = cast<T>(x, DataType::FLOAT32);
   }
-
-  std::vector<int64_t> x_dim = x_cast.shape();
-  int rank = x_dim.size();
+  int rank = x_cast.shape().size();
   DataLayout data_layout_ = common::StringToDataLayout(data_layout);
   int feature_axis;
   if (data_layout_ == DataLayout::kNCHW) {
@@ -390,68 +388,134 @@ std::tuple<Tensor, Tensor, Tensor, Tensor, Tensor, Tensor> batch_norm_decomp(
       reduce_axes.push_back(i);
     }
   }
-  std::vector<int64_t> stats_shape;
-  for (int i = 0; i < rank; ++i) {
-    if (find_value(reduce_axes, i) == false) {
-      stats_shape.push_back(x_dim[i]);
+  Tensor half = full_scalar<T>(-0.5, x_cast.dtype());
+  bool use_run_stat = (is_test && (!trainable_statistics)) || use_global_stats;
+
+  Tensor y, run_mean_, run_var_, batch_mean_, inv_std_, reserve_space;
+  if (has_dynamic_shape(x.shape())) {
+    Tensor x_dim = shape<T>(x_cast);
+    Tensor one_tensor = full<T>({1}, 1.0, x_dim.dtype());
+    std::vector<Tensor> stats_shape;
+    for (int i = 0; i < rank; i++) {
+      if (find_value(reduce_axes, i) == false) {
+        stats_shape.push_back(get_slice<T>(x_dim, i));
+      } else {
+        stats_shape.push_back(one_tensor);
+      }
+    }
+    Tensor stats_shape_tensor = concat<T>(stats_shape, 0);
+    Tensor x_hat;
+    Tensor batch_mean;
+    Tensor inv_std;
+
+    if (!use_run_stat) {
+      batch_mean = mean_decomp<T>(x_cast, reduce_axes, false);
+      auto temp = mean_decomp<T>(x_cast * x_cast, reduce_axes, false);
+      auto batch_var = temp - batch_mean * batch_mean;
+      auto eps = full_scalar<T>(epsilon, batch_var.dtype());
+      inv_std = elementwise_pow<T>((batch_var + eps), half);
+      if (data_layout_ == DataLayout::kNHWC) {
+        x_hat = (x_cast - batch_mean) * inv_std;
+      } else {
+        x_hat = (x_cast - backend::reshape<T>(batch_mean, stats_shape_tensor)) *
+                backend::reshape<T>(inv_std, stats_shape_tensor);
+      }
+      auto momentum_tensor = full_scalar<T>(momentum, run_mean.dtype());
+      auto momentum_sub_tensor =
+          full_scalar<T>(1. - momentum, run_mean.dtype());
+      run_mean_ = run_mean * momentum_tensor + batch_mean * momentum_sub_tensor;
+      run_var_ = run_var * momentum_tensor + batch_var * momentum_sub_tensor;
+      assign_out_<T>(run_mean_, run_mean);
+      assign_out_<T>(run_var_, run_var);
     } else {
-      stats_shape.push_back(1);
+      batch_mean =
+          backend::full_with_tensor<T>(shape<T>(run_mean), 0, run_mean.dtype());
+      auto batch_var =
+          backend::full_with_tensor<T>(shape<T>(run_var), 0, run_var.dtype());
+      auto eps = full_scalar<T>(epsilon, batch_var.dtype());
+      inv_std = elementwise_pow<T>((batch_var + eps), half);
+      if (data_layout_ == DataLayout::kNHWC) {
+        x_hat = (x_cast - run_mean) * elementwise_pow<T>((run_var + eps), half);
+      } else {
+        x_hat =
+            (x_cast - backend::reshape<T>(run_mean, stats_shape_tensor)) *
+            elementwise_pow<T>(
+                (backend::reshape<T>(run_var, stats_shape_tensor) + eps), half);
+      }
+      run_mean_ = assign<T>(run_mean);
+      run_var_ = assign<T>(run_var);
     }
-  }
 
-  Tensor half = full_scalar<T>(-0.5, x_cast.dtype());
-
-  bool use_run_stat = (is_test && (!trainable_statistics)) || use_global_stats;
-  Tensor x_hat;
-  Tensor batch_mean;
-  Tensor inv_std;
-  Tensor run_mean_;
-  Tensor run_var_;
-  if (!use_run_stat) {
-    batch_mean = mean_decomp<T>(x_cast, reduce_axes, false);
-    auto temp = mean_decomp<T>(x_cast * x_cast, reduce_axes, false);
-    auto batch_var = temp - batch_mean * batch_mean;
-    inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+    Tensor new_scale = scale ? scale.get() : full_scalar<T>(1, x_cast.dtype());
+    Tensor new_bias = bias ? bias.get() : full_scalar<T>(0, x_cast.dtype());
     if (data_layout_ == DataLayout::kNHWC) {
-      x_hat = (x_cast - batch_mean) * inv_std;
+      y = x_hat * new_scale + new_bias;
     } else {
-      x_hat = (x_cast - reshape<T>(batch_mean, stats_shape)) *
-              reshape<T>(inv_std, stats_shape);
+      y = x_hat * backend::reshape<T>(new_scale, stats_shape_tensor) +
+          backend::reshape<T>(new_bias, stats_shape_tensor);
+    }
+    batch_mean_ = assign<T>(batch_mean);
+    inv_std_ = assign<T>(inv_std);
+    if (need_cast) {
+      y = cast<T>(y, org_dtype);
     }
-    run_mean_ = run_mean * momentum + batch_mean * (1. - momentum);
-    run_var_ = run_var * momentum + batch_var * (1. - momentum);
-    assign_out_<T>(run_mean_, run_mean);
-    assign_out_<T>(run_var_, run_var);
   } else {
-    batch_mean = full<T>(run_mean.shape(), 0, run_mean.dtype());
-    auto batch_var = full<T>(run_var.shape(), 0, run_var.dtype());
-    inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+    std::vector<int64_t> x_dim = x_cast.shape();
+    std::vector<int64_t> stats_shape;
+    for (int i = 0; i < rank; ++i) {
+      if (find_value(reduce_axes, i) == false) {
+        stats_shape.push_back(x_dim[i]);
+      } else {
+        stats_shape.push_back(1);
+      }
+    }
+    Tensor x_hat;
+    Tensor batch_mean;
+    Tensor inv_std;
+    if (!use_run_stat) {
+      batch_mean = mean_decomp<T>(x_cast, reduce_axes, false);
+      auto temp = mean_decomp<T>(x_cast * x_cast, reduce_axes, false);
+      auto batch_var = temp - batch_mean * batch_mean;
+      inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+      if (data_layout_ == DataLayout::kNHWC) {
+        x_hat = (x_cast - batch_mean) * inv_std;
+      } else {
+        x_hat = (x_cast - reshape<T>(batch_mean, stats_shape)) *
+                reshape<T>(inv_std, stats_shape);
+      }
+      run_mean_ = run_mean * momentum + batch_mean * (1. - momentum);
+      run_var_ = run_var * momentum + batch_var * (1. - momentum);
+      assign_out_<T>(run_mean_, run_mean);
+      assign_out_<T>(run_var_, run_var);
+    } else {
+      batch_mean = full<T>(run_mean.shape(), 0, run_mean.dtype());
+      auto batch_var = full<T>(run_var.shape(), 0, run_var.dtype());
+      inv_std = elementwise_pow<T>((batch_var + epsilon), half);
+      if (data_layout_ == DataLayout::kNHWC) {
+        x_hat =
+            (x_cast - run_mean) * elementwise_pow<T>((run_var + epsilon), half);
+      } else {
+        x_hat = (x_cast - reshape<T>(run_mean, stats_shape)) *
+                elementwise_pow<T>((reshape<T>(run_var, stats_shape) + epsilon),
+                                   half);
+      }
+      run_mean_ = assign<T>(run_mean);
+      run_var_ = assign<T>(run_var);
+    }
+    Tensor new_scale = scale ? scale.get() : full_scalar<T>(1, x_cast.dtype());
+    Tensor new_bias = bias ? bias.get() : full_scalar<T>(0, x_cast.dtype());
     if (data_layout_ == DataLayout::kNHWC) {
-      x_hat =
-          (x_cast - run_mean) * elementwise_pow<T>((run_var + epsilon), half);
+      y = x_hat * new_scale + new_bias;
     } else {
-      x_hat = (x_cast - reshape<T>(run_mean, stats_shape)) *
-              elementwise_pow<T>((reshape<T>(run_var, stats_shape) + epsilon),
-                                 half);
-    }
-    run_mean_ = assign<T>(run_mean);
-    run_var_ = assign<T>(run_var);
-  }
-  Tensor y;
-  Tensor new_scale = scale ? scale.get() : full_scalar<T>(1, x_cast.dtype());
-  Tensor new_bias = bias ? bias.get() : full_scalar<T>(0, x_cast.dtype());
-  if (data_layout_ == DataLayout::kNHWC) {
-    y = x_hat * new_scale + new_bias;
-  } else {
-    y = x_hat * reshape<T>(new_scale, stats_shape) +
-        reshape<T>(new_bias, stats_shape);
-  }
-  Tensor reserve_space;
+      y = x_hat * reshape<T>(new_scale, stats_shape) +
+          reshape<T>(new_bias, stats_shape);
+    }
 
-  auto batch_mean_ = assign<T>(batch_mean);
-  auto inv_std_ = assign<T>(inv_std);
-  if (need_cast) {
-    y = cast<T>(y, org_dtype);
+    batch_mean_ = assign<T>(batch_mean);
+    inv_std_ = assign<T>(inv_std);
+    if (need_cast) {
+      y = cast<T>(y, org_dtype);
+    }
   }
   if (!use_run_stat) {
     return std::make_tuple(
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
