[Bug fixes]fix bce error when option pos_weight given (#65859)

LielinJiang · web-flow · commit 41a5220f54ba · 2024-07-23T16:11:34.000+08:00
* fix bce and bce grad with option pos_weight
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1340,15 +1340,19 @@ Tensor sigmoid_cross_entropy_with_logits_decomp(
   const Tensor zero = full_like_decomp<T>(x, 0, x.type(), x.place());
   const Tensor one = full_like_decomp<T>(x, 1, x.type(), x.place());
   Tensor pos_weight_tensor;
+  Tensor tmp_out;
   if (pos_weight) {
     pos_weight_tensor = pos_weight.get();
+    auto max_val = where<T>(x < zero, -x, zero);
+    auto term1 = (one - label) * x;
+    auto term2 = log<T>(exp<T>(-max_val) + exp<T>(-x - max_val));
+    tmp_out = term1 + pos_weight_tensor * (term2 + max_val);
   } else {
-    pos_weight_tensor = one;
+    auto term1 = where<T>(x > zero, x, zero);
+    auto term2 = x * label;
+    auto term3 = log<T>(one + exp<T>(-abs<T>(x)));
+    tmp_out = term1 - term2 + term3;
   }
-  auto term1 = where<T>(x > zero, x, zero);
-  auto term2 = x * label;
-  auto term3 = log<T>(one + exp<T>(-abs<T>(x)));
-  const Tensor tmp_out = term1 - term2 + term3 * pos_weight_tensor;
   const Tensor ignore_index_tensor =
       full_like_decomp<T>(x, ignore_index, label.type(), label.place());
   auto out = where<T>(label == ignore_index_tensor, zero, tmp_out);
diff --git a/paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_grad_kernel.cc b/paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_grad_kernel.cc
@@ -43,20 +43,30 @@ void SigmoidCrossEntropyWithLogitsGradKernel(
     T x = x_data[idx];
     T label = label_data[idx];
     T dout = dout_data[idx];
-    T pos_weight_idx = pos_weight_data == nullptr ? 1 : pos_weight_data[idx];
     if (static_cast<int>(label) == ignore_index) {
       dx_data[idx] = static_cast<T>(0.);
     } else {
-      T term1 = (x > 0) ? static_cast<T>(1) : static_cast<T>(0);
+      if (pos_weight_data == nullptr) {
+        T term1 = (x > 0) ? static_cast<T>(1) : static_cast<T>(0);
+        T e_x = std::exp(-std::abs(x));
+        T down = 1 + e_x;
+        T abs_grad = (x >= 0) ? static_cast<T>(1) : static_cast<T>(-1);
+        T up = -e_x * abs_grad;
+        T term3 = up / down;
 
-      T e_x = std::exp(-std::abs(x));
-      T down = 1 + e_x;
-      T abs_grad = (x >= 0) ? static_cast<T>(1) : static_cast<T>(-1);
-      T up = -e_x * abs_grad * pos_weight_idx;
-      T term3 = up / down;
-
-      T diff = term1 - label + term3;
-      dx_data[idx] = dout * diff;
+        T diff = term1 - label + term3;
+        dx_data[idx] = dout * diff;
+      } else {
+        T max_val = x < 0 ? -x : 0;
+        T term1 = (x < 0) ? static_cast<T>(-1) : static_cast<T>(0);
+        T down1 = std::exp(-max_val);
+        T down2 = std::exp(-x - max_val);
+        T term2 = down1 * (-term1) + down2 * (-1 - term1);
+        T term3 = (static_cast<T>(1.) - label);
+        T diff =
+            pos_weight_data[idx] * (term2 / (down1 + down2) + term1) + term3;
+        dx_data[idx] = dout * diff;
+      }
     }
   }
   if (normalize) {
diff --git a/paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_kernel.cc b/paddle/phi/kernels/cpu/sigmoid_cross_entropy_with_logits_kernel.cc
@@ -45,11 +45,17 @@ void SigmoidCrossEntropyWithLogitsKernel(
     if (static_cast<int>(label) == ignore_index) {
       out_data[idx] = static_cast<T>(0.);
     } else {
-      T pos_weight_idx = pos_weight_data == nullptr ? 1 : pos_weight_data[idx];
-      T term1 = (x > 0) ? x : 0;
-      T term2 = x * label;
-      T term3 = std::log(static_cast<T>(1) + std::exp(-std::abs(x)));
-      out_data[idx] = term1 - term2 + term3 * pos_weight_idx;
+      if (pos_weight_data == nullptr) {
+        T term1 = (x > 0) ? x : 0;
+        T term2 = x * label;
+        T term3 = std::log(static_cast<T>(1) + std::exp(-std::abs(x)));
+        out_data[idx] = term1 - term2 + term3;
+      } else {
+        T max_val = x < 0 ? -x : 0;
+        T term1 = (static_cast<T>(1.) - label) * x;
+        T term2 = std::log(std::exp(-max_val) + std::exp(-x - max_val));
+        out_data[idx] = term1 + pos_weight_data[idx] * (term2 + max_val);
+      }
     }
   }
 
diff --git a/paddle/phi/kernels/gpu/sigmoid_cross_entropy_with_logits_grad_kernel.cu b/paddle/phi/kernels/gpu/sigmoid_cross_entropy_with_logits_grad_kernel.cu
@@ -73,15 +73,15 @@ struct SigmoidBwdPosWeightFunctor {
       dx_data = static_cast<T>(0.);
       counts = 0;
     } else {
-      T term1 = (x > 0) ? static_cast<T>(1) : static_cast<T>(0);
-      T e_x = phi::funcs::real_exp(-abs(x));
-      T down = 1 + e_x;
-      T abs_grad = (x >= 0) ? static_cast<T>(1) : static_cast<T>(-1);
-      T up = -e_x * abs_grad * pos_weight;
-      T term3 = up / down;
-
-      T diff = term1 - label + term3;
+      T max_val = x < 0 ? -x : 0;
+      T term1 = (x < 0) ? static_cast<T>(-1) : static_cast<T>(0);
+      T down1 = phi::funcs::real_exp(-max_val);
+      T down2 = phi::funcs::real_exp(-x - max_val);
+      T term2 = down1 * (-term1) + down2 * (-1 - term1);
+      T term3 = (static_cast<T>(1.) - label);
+      T diff = pos_weight * (term2 / (down1 + down2) + term1) + term3;
       dx_data = dout * diff;
+
       counts = 1;
     }
     phi::Array<T, 2> outs;
diff --git a/paddle/phi/kernels/gpu/sigmoid_cross_entropy_with_logits_kernel.cu b/paddle/phi/kernels/gpu/sigmoid_cross_entropy_with_logits_kernel.cu
@@ -72,14 +72,12 @@ struct SigmoidFwdPosWeightFunctor {
       out_data = static_cast<T>(0.);
       counts = 0;
     } else {
-      T term1 = (x > 0) ? x : 0;
-      T term2 = x * label;
-      T term3 =
-          phi::funcs::real_log(static_cast<T>(1) +
-                               phi::funcs::real_exp(static_cast<T>(-abs(x)))) *
-          pos_weight;
+      T max_val = x < 0 ? -x : 0;
+      T term1 = (static_cast<T>(1.) - label) * x;
+      T term2 = phi::funcs::real_log(phi::funcs::real_exp(-max_val) +
+                                     phi::funcs::real_exp(-x - max_val));
+      out_data = term1 + pos_weight * (term2 + max_val);
 
-      out_data = term1 - term2 + term3;
       counts = 1;
     }
     phi::Array<T, 2> outs;
diff --git a/test/deprecated/legacy_test/test_sigmoid_cross_entropy_with_logits_op.py b/test/deprecated/legacy_test/test_sigmoid_cross_entropy_with_logits_op.py
@@ -173,10 +173,12 @@ def setUp(self):
 
         # Fw Pass is implemented as elementwise sigmoid followed by
         # elementwise logistic loss
-        term1 = np.maximum(self.inputs['X'], 0)
-        term2 = self.inputs['X'] * self.inputs['Label']
-        term3 = np.log(1 + np.exp(-1 * np.abs(self.inputs['X']))) * pos_weight
-        self.outputs = {'Out': term1 - term2 + term3}
+        max_val = np.clip(-self.inputs['X'], 0, np.finfo(np.float64).max)
+        term1 = (1 - label) * self.inputs['X']
+        term2 = np.log(np.exp(-max_val) + np.exp(-self.inputs['X'] - max_val))
+        out = term1 + pos_weight * (term2 + max_val)
+
+        self.outputs = {'Out': out}
 
     def test_check_output(self):
         self.check_output(check_pir=True, check_prim_pir=True)
diff --git a/test/legacy_test/test_sigmoid_cross_entropy_with_logits_grad_with_auto_grad.py b/test/legacy_test/test_sigmoid_cross_entropy_with_logits_grad_with_auto_grad.py
@@ -56,12 +56,13 @@ def fn_ref(x, label, weight):
 
         def fn_comp(x, label, weight):
             zeros = paddle.full((self.batch_size, self.num_classes), 0.0)
-            t1 = paddle.where(x > 0, x, zeros)
-            t2 = x * label
-            t3 = paddle.log(1 + paddle.exp(-paddle.abs(x)))
-            t4 = t1 - t2 + t3 * weight
-            t5 = paddle.full((self.batch_size, self.num_classes), -100.0)
-            out = paddle.where(label == t5, zeros, t4)
+
+            max_val = paddle.where(x < zeros, -x, zeros)
+            t1 = (1 - label) * x
+            t2 = paddle.log((-max_val).exp() + (-x - max_val).exp())
+            t3 = t1 + weight * (t2 + max_val)
+            t4 = paddle.full((self.batch_size, self.num_classes), -100.0)
+            out = paddle.where(label == t4, zeros, t3)
             loss = out.sum()
             loss.backward()
             return out, x.grad
