[Prim][PIR] Forward decomposite the log_loss op (#65968)

zeroRains · web-flow · commit 802cc9275430 · 2024-07-15T16:56:55.000+08:00
* forward decomposite the log_loss op

* add the dynamic shape surport

* fix the tol
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -41,6 +41,7 @@
     "instance_norm",
     "layer_norm",
     "leaky_relu",
+    "log_loss",
     "log_softmax",
     "mean",
     "mean_all",
@@ -86,6 +87,7 @@
     "instance_norm",
     "layer_norm",
     "leaky_relu",
+    "log_loss",
     "log_softmax",
     "mean",
     "mean_all",
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1514,6 +1514,17 @@ Tensor elu_decomp(const Tensor& x, const float alpha) {
   }
 }
 
+template <typename T>
+Tensor log_loss_decomp(const Tensor& input,
+                       const Tensor& label,
+                       float epsilon) {
+  Tensor ones = full_scalar<T>(1.0, input.dtype());
+  Tensor eps = full_scalar<T>(epsilon);
+  Tensor term1 = -label * log<T>(input + eps);
+  Tensor term2 = (ones - label) * log<T>(ones - input + eps);
+  return term1 - term2;
+}
+
 }  // namespace details
 
 }  // namespace primitive
diff --git a/test/legacy_test/test_log_loss_op.py b/test/legacy_test/test_log_loss_op.py
@@ -28,6 +28,8 @@ class TestLogLossOp(OpTest):
     def setUp(self):
         self.op_type = 'log_loss'
         self.python_api = F.log_loss
+        self.prim_op_type = "comp"
+        self.public_python_api = F.log_loss
         samples_num = 100
 
         x = np.random.random((samples_num, 1)).astype("float32")
@@ -46,11 +48,15 @@ def setUp(self):
         self.outputs = {'Loss': loss}
 
     def test_check_output(self):
-        self.check_output(check_pir=True)
+        self.check_output(check_pir=True, check_prim_pir=True)
 
     def test_check_grad(self):
         self.check_grad(
-            ['Predicted'], 'Loss', max_relative_error=0.03, check_pir=True
+            ['Predicted'],
+            'Loss',
+            max_relative_error=0.03,
+            check_pir=True,
+            check_prim_pir=True,
         )
 
 
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
@@ -180,6 +180,10 @@ def unbind_net(x):
     return paddle.unbind(x)
 
 
+def log_loss_net(inputs, labels):
+    return paddle.nn.functional.log_loss(inputs, labels)
+
+
 class TestPrimBase(unittest.TestCase):
     def setUp(self):
         np.random.seed(2023)
@@ -645,5 +649,39 @@ def setUp(self):
         self.enable_cinn = False
 
 
+class TestPrimLogLoss1(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.x_shape = [400, 1]
+        self.y_shape = [400, 1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, None]
+        self.init_y_shape = [None, None]
+        self.x = np.random.random(self.x_shape).astype(self.dtype_x)
+        self.y = np.random.randint(0, 2, self.y_shape).astype(self.dtype_y)
+        self.net = log_loss_net
+        self.necessary_ops = "pd_op.log_loss"
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
+class TestPrimLogLoss2(TestPrimTwo):
+    def setUp(self):
+        np.random.seed(2023)
+        self.x_shape = [400, 1]
+        self.y_shape = [400, 1]
+        self.dtype_x = "float32"
+        self.dtype_y = "float32"
+        self.init_x_shape = [None, 1]
+        self.init_y_shape = [None, 1]
+        self.x = np.random.random(self.x_shape).astype(self.dtype_x)
+        self.y = np.random.randint(0, 2, self.y_shape).astype(self.dtype_y)
+        self.net = log_loss_net
+        self.necessary_ops = "pd_op.log_loss"
+        self.enable_cinn = False
+        self.tol = 1e-5
+
+
 if __name__ == "__main__":
     unittest.main()
