Fix the slow running speed of kl_div when option 'reduction' is set (#…

…37283)

* Fix the slow running speed of kl_div when option reduction is set

* fix unittest coverage

python/paddle/fluid/tests/unittests/test_kldiv_loss_op.py

@@ -112,7 +112,9 @@ def test_kl_loss_static_api(self):
         input = paddle.fluid.data(name='input', shape=[5, 20])
         label = paddle.fluid.data(name='label', shape=[5, 20])
 
-        pred_loss = paddle.nn.functional.kl_div(input, label)
+        paddle.nn.functional.kl_div(input, label)
+        paddle.nn.functional.kl_div(input, label, 'sum')
+        paddle.nn.functional.kl_div(input, label, 'batchmean')
 
 
 class TestKLDivLossTypePromotion(unittest.TestCase):

python/paddle/nn/functional/loss.py

@@ -903,7 +903,15 @@ def kl_div(input, label, reduction='mean', name=None):
         label = paddle.cast(label, 'float64')
 
     if paddle.in_dynamic_mode():
-        out = _C_ops.kldiv_loss(input, label, 'reduction', reduction)
+        out = _C_ops.kldiv_loss(input, label, 'reduction', 'none')
+        if reduction == 'mean':
+            out = paddle.mean(out)
+        elif reduction == 'sum':
+            out = paddle.sum(out)
+        elif reduction == 'batchmean':
+            if len(input.shape) > 0:
+                batch_size = input.shape[0]
+                out = paddle.sum(out) / batch_size
         return out
 
     helper = LayerHelper('kl_div', **locals())
@@ -920,7 +928,15 @@ def kl_div(input, label, reduction='mean', name=None):
         inputs={'X': input,
                 'Target': label},
         outputs={'Loss': loss},
-        attrs={'reduction': reduction})
+        attrs={'reduction': 'none'})
+
+    if reduction == 'mean':
+        loss = paddle.mean(loss)
+    elif reduction == 'sum':
+        loss = paddle.sum(loss)
+    elif reduction == 'batchmean':
+        batch_size = paddle.shape(input)[0]
+        loss = paddle.sum(loss) / batch_size
     return loss