support batch comupute for chamferloss

Author: HydrogenSulfate

ppsci/loss/chamfer.py

@@ -40,14 +40,17 @@ class ChamferLoss(base.Loss):
     Examples:
         >>> import paddle
         >>> from ppsci.loss import ChamferLoss
-        >>> output_dict = {"s1": paddle.to_tensor([[.1, .2, .3], [.4, .5, .6]])}
-        >>> label_dict  = {"s1": paddle.to_tensor([[.4, .5, .6], [.1, .2, .3]])}
+        >>> _ = paddle.seed(42)
+        >>> batch_point_cloud1 = paddle.rand([2, 100, 3])
+        >>> batch_point_cloud2 = paddle.rand([2, 50, 3])
+        >>> output_dict = {"s1": batch_point_cloud1}
+        >>> label_dict  = {"s1": batch_point_cloud2}
         >>> weight = {"s1": 0.8}
         >>> loss = ChamferLoss(weight=weight)
         >>> result = loss(output_dict, label_dict)
         >>> print(result)
         Tensor(shape=[], dtype=float32, place=Place(gpu:0), stop_gradient=True,
-               0.)
+               0.04415882)
     """
 
     def __init__(
@@ -61,16 +64,16 @@ def forward(self, output_dict, label_dict, weight_dict=None):
         for key in label_dict:
             s1 = output_dict[key]
             s2 = label_dict[key]
-            N1, N2 = s1.shape[0], s2.shape[0]
+            N1, N2 = s1.shape[1], s2.shape[1]
 
             # [N1, N2, 3]
-            s1_expand = paddle.expand(s1.reshape([N1, 1, 3]), shape=[N1, N2, 3])
+            s1_expand = paddle.expand(s1.reshape([-1, N1, 1, 3]), shape=[-1, N1, N2, 3])
             # [N1, N2, 3]
-            s2_expand = paddle.expand(s2.reshape([1, N2, 3]), shape=[N1, N2, 3])
+            s2_expand = paddle.expand(s2.reshape([-1, 1, N2, 3]), shape=[-1, N1, N2, 3])
 
-            dis = ((s1_expand - s2_expand) ** 2).sum(axis=2)  # [N1, N2]
-            loss_s12 = dis.min(axis=1)  # [N1]
-            loss_s21 = dis.min(axis=0)  # [N2]
+            dis = ((s1_expand - s2_expand) ** 2).sum(axis=3)  # [B, N1, N2]
+            loss_s12 = dis.min(axis=2)  # [B, N1]
+            loss_s21 = dis.min(axis=1)  # [B, N2]
             loss = loss_s12.mean() + loss_s21.mean()
 
             if weight_dict and key in weight_dict:

test/loss/chamfer.py

@@ -11,16 +11,16 @@ def test_chamfer_loss():
     """Test for chamfer distance loss."""
     N1 = 100
     N2 = 50
-    output_dict = {"s1": paddle.randn([N1, 3])}
-    label_dict = {"s1": paddle.randn([N2, 3])}
+    output_dict = {"s1": paddle.randn([1, N1, 3])}
+    label_dict = {"s1": paddle.randn([1, N2, 3])}
     chamfer_loss = loss.ChamferLoss()
     result = chamfer_loss(output_dict, label_dict)
 
     loss_cd_s1 = 0.0
     for i in range(N1):
         min_i = None
         for j in range(N2):
-            disij = ((output_dict["s1"][i] - label_dict["s1"][j]) ** 2).sum()
+            disij = ((output_dict["s1"][0, i] - label_dict["s1"][0, j]) ** 2).sum()
             if min_i is None or disij < min_i:
                 min_i = disij
         loss_cd_s1 += min_i
@@ -30,7 +30,7 @@ def test_chamfer_loss():
     for j in range(N2):
         min_j = None
         for i in range(N1):
-            disij = ((output_dict["s1"][i] - label_dict["s1"][j]) ** 2).sum()
+            disij = ((output_dict["s1"][0, i] - label_dict["s1"][0, j]) ** 2).sum()
             if min_j is None or disij < min_j:
                 min_j = disij
         loss_cd_s2 += min_j