API Improvement for paddle.nn.initializer.TruncatedNormal 易用性提升 (#64110)

NKNaN · web-flow · commit 657bb7e8a82e · 2024-05-16T14:27:43.000+08:00
* update truncated_gaussian_random kernel

* update docs

* fix error

* fix opp_version

* fix build error

* fix test

* add note in en docs and fix xpu test

* fix xpu test

* fix xpu test

* update

* update fluid

* rm fluid operator

* fix xpu test

* fix xpu test

* update note
diff --git a/paddle/fluid/distributed/ps/table/depends/initializers.h b/paddle/fluid/distributed/ps/table/depends/initializers.h
@@ -117,20 +117,22 @@ class TruncatedGaussianInitializer : public Initializer {
     seed_ = static_cast<unsigned int>(std::stoi(attrs[1]));
     mean_ = std::stof(attrs[2]);
     std_ = std::stof(attrs[3]);
+    a_ = std::stof(attrs[4]);
+    b_ = std::stof(attrs[5]);
 
     std::uniform_real_distribution<float> dist_(
         std::numeric_limits<float>::min(), 1.0);
     random_engine_ = phi::GetCPURandomEngine(seed_);
   }
 
   float GetValue() override {
-    TruncatedNormal<float> truncated_normal(mean_, std_);
+    TruncatedNormal<float> truncated_normal(mean_, std_, a_, b_);
     float value = truncated_normal(dist_(*random_engine_));
     return value;
   }
 
   void GetValue(float *value, int numel) {
-    TruncatedNormal<float> truncated_normal(mean_, std_);
+    TruncatedNormal<float> truncated_normal(mean_, std_, a_, b_);
     for (int x = 0; x < numel; ++x) {
       value[x] = truncated_normal(dist_(*random_engine_));
     }
@@ -139,6 +141,8 @@ class TruncatedGaussianInitializer : public Initializer {
  private:
   float std_;
   float mean_;
+  float a_;
+  float b_;
 
   std::shared_ptr<std::mt19937_64> random_engine_;
   std::uniform_real_distribution<float> dist_;
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1868,14 +1868,14 @@
 
 # python API: paddle.nn.initializer.TruncatedNormal
 - op : truncated_gaussian_random
-  args : (int[] shape, float mean, float std, int seed, DataType dtype=DataType::FLOAT32, Place place={})
+  args : (int[] shape, float mean, float std, int seed, float a, float b, DataType dtype=DataType::FLOAT32, Place place={})
   output : Tensor(out)
   infer_meta :
     func : TruncatedGaussianRandomInferMeta
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
   kernel :
     func : truncated_gaussian_random
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
     backend : place
     data_type : dtype
   traits : pir::SideEffectTrait
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -1259,14 +1259,14 @@
 
 # python API: paddle.nn.initializer.TruncatedNormal
 - op : truncated_gaussian_random
-  args : (int[] shape, float mean, float std, int seed, DataType dtype=DataType::FLOAT32, Place place={})
+  args : (int[] shape, float mean, float std, int seed, float a, float b, DataType dtype=DataType::FLOAT32, Place place={})
   output : Tensor(out)
   infer_meta :
     func : TruncatedGaussianRandomInferMeta
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
   kernel :
     func : truncated_gaussian_random
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
     backend : place
     data_type : dtype
 
diff --git a/paddle/phi/api/yaml/op_version.yaml b/paddle/phi/api/yaml/op_version.yaml
@@ -490,6 +490,17 @@
         - delete_attr : dim2
           comment : The attribute 'dim2' is not recommend according to the specification 2.0.
 
+- op : truncated_gaussian_random
+  version :
+    - checkpoint : Upgrade truncated_gaussian_random add 2 new attribute [a, b]
+      action :
+        - add_attr : a
+          comment : The minimum cutoff value.
+          default: -2.0
+        - add_attr : b
+          comment : The maximum cutoff value.
+          default: 2.0
+
 - op : unique_consecutive
   version :
     - checkpoint : Upgrade unique_consecutive, add 2 outputs [Indices, Counts] and 3 attribute [return_inverse, return_counts, axis].
diff --git a/paddle/phi/api/yaml/static_ops.yaml b/paddle/phi/api/yaml/static_ops.yaml
@@ -699,14 +699,14 @@
     data_type : dtype
 
 - op : truncated_gaussian_random
-  args : (int[] shape, float mean = .0f, float std = 1.0f, int seed = 0, DataType dtype=DataType::FLOAT32)
+  args : (int[] shape, float mean = .0f, float std = 1.0f, int seed = 0, float a = -2.0f, float b = 2.0f, DataType dtype=DataType::FLOAT32)
   output : Tensor(out)
   infer_meta :
     func : TruncatedGaussianRandomInferMeta
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
   kernel :
     func : truncated_gaussian_random
-    param : [shape, mean, std, seed, dtype]
+    param : [shape, mean, std, seed, a, b, dtype]
     data_type : dtype
 
 - op : uniform
diff --git a/paddle/phi/infermeta/nullary.cc b/paddle/phi/infermeta/nullary.cc
@@ -263,6 +263,8 @@ void TruncatedGaussianRandomInferMeta(const std::vector<int>& shape,
                                       float mean,
                                       float std,
                                       int seed,
+                                      float a,
+                                      float b,
                                       DataType dtype,
                                       MetaTensor* out) {
   auto out_dims = common::make_ddim(shape);
diff --git a/paddle/phi/infermeta/nullary.h b/paddle/phi/infermeta/nullary.h
@@ -109,6 +109,8 @@ void TruncatedGaussianRandomInferMeta(const std::vector<int>& shape,
                                       float mean,
                                       float std,
                                       int seed,
+                                      float a,
+                                      float b,
                                       DataType dtype,
                                       MetaTensor* out);
 
diff --git a/paddle/phi/kernels/cpu/truncated_gaussian_random_kernel.cc b/paddle/phi/kernels/cpu/truncated_gaussian_random_kernel.cc
@@ -30,6 +30,8 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
                                    float mean,
                                    float std,
                                    int seed,
+                                   float a,
+                                   float b,
                                    DataType dtype,
                                    DenseTensor* out) {
   auto tensor = out;
@@ -38,7 +40,7 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
 
   std::uniform_real_distribution<T> dist(std::numeric_limits<float>::min(),
                                          1.0);
-  TruncatedNormal<T> truncated_normal(mean, std);
+  TruncatedNormal<T> truncated_normal(mean, std, a, b);
   int64_t size = tensor->numel();
 
   std::shared_ptr<std::mt19937_64> engine;
diff --git a/paddle/phi/kernels/funcs/truncated_normal.h b/paddle/phi/kernels/funcs/truncated_normal.h
@@ -130,19 +130,20 @@ T Erfinv(T x) {
 
 template <typename T>
 struct TruncatedNormal {
-  T mean, std;
+  T mean, std, a, b;
   T a_normal_cdf;
   T b_normal_cdf;
-  TruncatedNormal(T mean, T std) : mean(mean), std(std) {
+  TruncatedNormal(T mean, T std, T a, T b) : mean(mean), std(std), a(a), b(b) {
     auto normal_cdf = [](T x) {
       return (1.0 + std::erf(x / std::sqrt(2.0))) / 2.0;
     };
-    a_normal_cdf = normal_cdf(-2.0);
-    b_normal_cdf = normal_cdf(2.0);
+    a_normal_cdf = normal_cdf((a - mean) / std);
+    b_normal_cdf = normal_cdf((b - mean) / std);
   }
 
   T operator()(T value) const {
     auto p = a_normal_cdf + (b_normal_cdf - a_normal_cdf) * value;
-    return std::sqrt(2.0) * Erfinv(2 * p - 1) * std + mean;
+    T ret = std::sqrt(2.0) * Erfinv(2 * p - 1) * std + mean;
+    return std::clamp(ret, a, b);
   }
 };
diff --git a/paddle/phi/kernels/gpu/truncated_gaussian_random_kernel.cu b/paddle/phi/kernels/gpu/truncated_gaussian_random_kernel.cu
@@ -29,16 +29,17 @@ namespace phi {
 
 template <typename T>
 struct GPUTruncatedNormal {
-  T mean, std;
+  T mean, std, a, b;
   T a_normal_cdf;
   T b_normal_cdf;
   unsigned int seed;
   T numeric_min;
 
-  __host__ __device__ GPUTruncatedNormal(T mean, T std, T numeric_min, int seed)
-      : mean(mean), std(std), seed(seed), numeric_min(numeric_min) {
-    a_normal_cdf = (1.0 + erff(-2.0 / sqrtf(2.0))) / 2.0;
-    b_normal_cdf = (1.0 + erff(2.0 / sqrtf(2.0))) / 2.0;
+  __host__ __device__
+  GPUTruncatedNormal(T mean, T std, T numeric_min, int seed, T a, T b)
+      : mean(mean), std(std), seed(seed), numeric_min(numeric_min), a(a), b(b) {
+    a_normal_cdf = (1.0 + erff((a - mean) / std / sqrtf(2.0))) / 2.0;
+    b_normal_cdf = (1.0 + erff((b - mean) / std / sqrtf(2.0))) / 2.0;
   }
 
   __host__ __device__ T operator()(const unsigned int n) const {
@@ -48,28 +49,31 @@ struct GPUTruncatedNormal {
     rng.discard(n);
     T value = dist(rng);
     auto p = a_normal_cdf + (b_normal_cdf - a_normal_cdf) * value;
-    return std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    T ret = std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    return std::clamp(ret, a, b);
   }
 };
 
 template <typename T>
 struct TruncatedNormalOffset {
-  T mean, std;
+  T mean, std, a, b;
   T a_normal_cdf;
   T b_normal_cdf;
   unsigned int seed;
   T numeric_min;
   int offset_;
 
-  __host__ __device__
-  TruncatedNormalOffset(T mean, T std, T numeric_min, int seed, int offset)
+  __host__ __device__ TruncatedNormalOffset(
+      T mean, T std, T numeric_min, int seed, int offset, T a, T b)
       : mean(mean),
         std(std),
         seed(seed),
         numeric_min(numeric_min),
-        offset_(offset) {
-    a_normal_cdf = (1.0 + erff(-2.0 / sqrtf(2.0))) / 2.0;
-    b_normal_cdf = (1.0 + erff(2.0 / sqrtf(2.0))) / 2.0;
+        offset_(offset),
+        a(a),
+        b(b) {
+    a_normal_cdf = (1.0 + erff((a - mean) / std / sqrtf(2.0))) / 2.0;
+    b_normal_cdf = (1.0 + erff((b - mean) / std / sqrtf(2.0))) / 2.0;
   }
 
   __host__ __device__ T operator()(const unsigned int n) const {
@@ -79,7 +83,8 @@ struct TruncatedNormalOffset {
     rng.discard(n + offset_);
     T value = dist(rng);
     auto p = a_normal_cdf + (b_normal_cdf - a_normal_cdf) * value;
-    return std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    T ret = std::sqrt(2.0) * erfinvf(2 * p - 1) * std + mean;
+    return std::clamp(ret, a, b);
   }
 };
 
@@ -89,6 +94,8 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
                                    float mean,
                                    float std,
                                    int seed,
+                                   float a,
+                                   float b,
                                    DataType dtype,
                                    DenseTensor* out) {
   T* data = dev_ctx.template Alloc<T>(out);
@@ -102,19 +109,24 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
     auto seed_offset = gen_cuda->IncrementOffset(1);
     uint64_t seed = seed_offset.first;
     uint64_t offset = seed_offset.second;
-    thrust::transform(
-        index_sequence_begin,
-        index_sequence_begin + size,
-        thrust::device_ptr<T>(data),
-        TruncatedNormalOffset<T>(
-            mean, std, std::numeric_limits<T>::min(), seed, size * offset));
+    thrust::transform(index_sequence_begin,
+                      index_sequence_begin + size,
+                      thrust::device_ptr<T>(data),
+                      TruncatedNormalOffset<T>(mean,
+                                               std,
+                                               std::numeric_limits<T>::min(),
+                                               seed,
+                                               size * offset,
+                                               a,
+                                               b));
   } else {
     // use OP seed
     thrust::transform(
         index_sequence_begin,
         index_sequence_begin + size,
         thrust::device_ptr<T>(data),
-        GPUTruncatedNormal<T>(mean, std, std::numeric_limits<T>::min(), seed));
+        GPUTruncatedNormal<T>(
+            mean, std, std::numeric_limits<T>::min(), seed, a, b));
   }
 }
 
diff --git a/paddle/phi/kernels/truncated_gaussian_random_kernel.h b/paddle/phi/kernels/truncated_gaussian_random_kernel.h
@@ -25,6 +25,8 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
                                    float mean,
                                    float std,
                                    int seed,
+                                   float a,
+                                   float b,
                                    DataType dtype,
                                    DenseTensor* out);
 
diff --git a/paddle/phi/kernels/xpu/truncated_gaussian_random_kernel.cc b/paddle/phi/kernels/xpu/truncated_gaussian_random_kernel.cc
@@ -30,13 +30,15 @@ void TruncatedGaussianRandomKernel(const Context& dev_ctx,
                                    float mean,
                                    float std,
                                    int seed,
+                                   float a,
+                                   float b,
                                    DataType dtype,
                                    DenseTensor* out) {
   T* data = dev_ctx.template Alloc<T>(out);
 
   std::uniform_real_distribution<T> dist(std::numeric_limits<float>::min(),
                                          1.0);
-  TruncatedNormal<T> truncated_normal(mean, std);
+  TruncatedNormal<T> truncated_normal(mean, std, a, b);
   int64_t size = out->numel();
   std::unique_ptr<T[]> data_cpu(new T[size]);
 
diff --git a/python/paddle/nn/initializer/normal.py b/python/paddle/nn/initializer/normal.py
@@ -165,21 +165,31 @@ def __init__(self, mean=0.0, std=1.0, name=None):
 class TruncatedNormalInitializer(Initializer):
     """Implements the Random TruncatedNormal(Gaussian) distribution initializer
 
+    Note:
+        It is better to set `a <= mean <= b`.
+        If `mean < a - 2*std` or `mean > b + 2*std`, the distribution of values may be incorrect.
+
     Args:
         loc (float, optional): Mean of the normal distribution. Default is :math:`0.0`.
         scale (float, optional): Standard deviation of the normal distribution. Default is :math:`1.0`.
         seed (int, optional): random seed. Default is 0.
+        a (float, optional): The minimum cutoff value. Default is -2.0.
+        b (float, optional): The maximum cutoff value. Default is 2.0.
 
     """
 
-    def __init__(self, loc=0.0, scale=1.0, seed=0):
+    def __init__(self, loc=0.0, scale=1.0, seed=0, a=-2.0, b=2.0):
         assert loc is not None
         assert scale is not None
         assert seed is not None
+        assert a is not None
+        assert b is not None
         super().__init__()
         self._mean = loc
         self._std_dev = scale
         self._seed = seed
+        self._a = a
+        self._b = b
 
     def forward(self, var, block=None):
         """Initialize the input tensor with TruncatedNormal distribution.
@@ -222,6 +232,8 @@ def forward(self, var, block=None):
                 self._mean,
                 self._std_dev,
                 self._seed,
+                self._a,
+                self._b,
                 out_dtype,
                 _current_expected_place(),
             )
@@ -245,6 +257,8 @@ def forward(self, var, block=None):
                     "mean": self._mean,
                     "std": self._std_dev,
                     "seed": self._seed,
+                    "a": self._a,
+                    "b": self._b,
                 },
                 stop_gradient=True,
             )
@@ -266,9 +280,15 @@ def forward(self, var, block=None):
 class TruncatedNormal(TruncatedNormalInitializer):
     """The truncated normal distribution (Gaussian distribution) initializer.
 
+    Note:
+        It is better to set `a <= mean <= b`.
+        If `mean < a - 2*std` or `mean > b + 2*std`, the distribution of values may be incorrect.
+
     Args:
         mean (float, optional): Mean of the normal distribution. Default is :math:`0.0`.
         std (float, optional): Standard deviation of the normal distribution. Default is :math:`1.0`.
+        a (float, optional): The minimum cutoff value. Default is -2.0.
+        b (float, optional): The maximum cutoff value. Default is 2.0.
         name (str, optional): For details, please refer to :ref:`api_guide_Name`. Generally, no setting is required. Default: None.
 
     Returns:
@@ -305,7 +325,9 @@ class TruncatedNormal(TruncatedNormalInitializer):
              [[-0.11380529 -3.0696259 ]]])
     """
 
-    def __init__(self, mean=0.0, std=1.0, name=None):
+    def __init__(self, mean=0.0, std=1.0, a=-2.0, b=2.0, name=None):
         assert mean is not None, 'mean should not be None'
         assert std is not None, 'std should not be None'
-        super().__init__(loc=mean, scale=std, seed=0)
+        assert a is not None, 'a should not be None'
+        assert b is not None, 'b should not be None'
+        super().__init__(loc=mean, scale=std, seed=0, a=a, b=b)
diff --git a/test/deprecated/legacy_test/test_initializer_nn.py b/test/deprecated/legacy_test/test_initializer_nn.py
diff --git a/test/deprecated/legacy_test/test_truncated_gaussian_random_op.py b/test/deprecated/legacy_test/test_truncated_gaussian_random_op.py
diff --git a/test/xpu/test_truncated_gaussian_random_op_xpu.py b/test/xpu/test_truncated_gaussian_random_op_xpu.py
