change input dimension from (B, H, W, C) to (B, C, H, W)

Author: pri1311

test/test_ops.py

@@ -1775,18 +1775,22 @@ def get_reduction_method(self, reduction):
 
     @pytest.mark.parametrize("device", cpu_and_gpu())
     def test_dice_loss(self, device):
-        input = torch.tensor([[[0.9409, 0.9220], [0.9524, 0.1094]],
-                              [[0.6802, 0.7949], [0.9570, 0.1499]],
-                              [[0.3298, 0.4401], [0.1094, 0.7536]],
-                              [[0.3340, 0.9895], [0.9563, 0.5045]]], device=device)
+        input = torch.tensor([[[0.9409, 0.9524],
+                               [0.9220, 0.1094]],
+                              [[0.6802, 0.9570],
+                               [0.7949, 0.1499]],
+                              [[0.3298, 0.1094],
+                               [0.4401, 0.7536]],
+                              [[0.3340, 0.9563],
+                               [0.9895, 0.5045]]], device=device)
         labels = torch.tensor([[[0, 1], [1, 0]],
                                [[1, 0], [0, 1]],
-                               [[1, 0], [1, 0]],
+                               [[1, 1], [0, 0]],
                                [[1, 0], [0, 1]]], device=device)
         expected = torch.tensor([0.4028, 0.6101, 0.5916, 0.6347], device=device)
         torch.testing.assert_allclose(ops.dice_loss(input, labels, eps=0), expected)
 
-    @pytest.mark.parametrize("shape", ((16, 4, 4, 2), (32, 2), (32, 4, 4, 4, 2)))
+    @pytest.mark.parametrize("shape", ((16, 2, 4, 4), (16, 4, 4, 4), (32, 2), (32, 2, 4, 4, 4)))
     @pytest.mark.parametrize("reduction", ["none", "mean", "sum"])
     @pytest.mark.parametrize("device", cpu_and_gpu())
     def test_dice_loss_one(self, shape, reduction, device):
@@ -1800,19 +1804,19 @@ def test_dice_loss_one(self, shape, reduction, device):
 
     @pytest.mark.parametrize("device", cpu_and_gpu())
     def test_dice_loss_all_zeros(self, device):
-        shape = (16, 4, 4, 2)
+        shape = (16, 2, 4, 4)
         input_zeros = torch.zeros(shape, device=device)
-        input_zeros[:, :, :, 0] = 1.0
-        input_zeros[:, :, :, 1] = 0.0
+        input_zeros[:, 0, :, :] = 1.0
+        input_zeros[:, 1, :, :] = 0.0
         label_zeros = torch.zeros(shape, device=device)
         label_zeros.copy_(input_zeros)
-        input_zeros[:, :, :, 0] = 100.0
+        input_zeros[:, 0, :, :] = 100.0
         expected = torch.zeros(16, device=device)
         torch.testing.assert_close(ops.dice_loss(input_zeros, label_zeros), expected)
 
     @pytest.mark.parametrize("device", cpu_and_gpu())
     def test_gradcheck(self, device):
-        shape = (16, 4, 4, 2)
+        shape = (16, 2, 4, 4)
         input_ones = torch.ones(shape, device=device, requires_grad=True)
         label_zeros = torch.zeros(shape, device=device, requires_grad=True)
         assert gradcheck(ops.dice_loss, (input_ones, label_zeros), eps=1e-2, atol=1e-2, raise_exception=True, fast_mode=True)

torchvision/ops/dice_loss.py

@@ -24,11 +24,11 @@ def dice_loss(inputs: torch.Tensor, targets: torch.Tensor, reduction: str = "non
         \text{loss}(x, class) = 1 - \text{Dice}(x, class)
 
     Args:
-        inputs: (Tensor): A float tensor with rank >= 2 and shape (B, N1, .... NK, C)
+        inputs: (Tensor): A float tensor with rank >= 2 and shape (B, C, N1, .... NK)
                 where B is the Batch Size and C is the number of classes.
                 The predictions for each example.
         targets: (Tensor): A one-hot tensor with the same shape as inputs.
-                The first dimension is the batch size and last dimension is the
+                The first dimension is the batch size and the second dimension is the
                 number of classes.
         eps: (float, optional): Scalar to enforce numerical stability.
         reduction (string, optional): ``'none'`` | ``'mean'`` | ``'sum'``
@@ -44,13 +44,13 @@ def dice_loss(inputs: torch.Tensor, targets: torch.Tensor, reduction: str = "non
         _log_api_usage_once(dice_loss)
 
     # compute softmax over the classes axis
-    p = F.softmax(inputs, dim=-1)
+    p = F.softmax(inputs, dim=1)
     p = p.flatten(start_dim=1)
 
     targets = targets.flatten(start_dim=1)
 
-    intersection = torch.sum(p * targets, dim=-1)
-    cardinality = torch.sum(p + targets, dim=-1)
+    intersection = torch.sum(p * targets, dim=1)
+    cardinality = torch.sum(p + targets, dim=1)
 
     dice_score = 2.0 * intersection / (cardinality + eps)