add reduce option

Author: ir2718

src/pytorch_metric_learning/losses/smooth_ap.py

@@ -80,19 +80,20 @@ def compute_loss(self, embeddings, labels, indices_tuple, ref_emb, ref_labels):
         )
         sims_pos_ranks = torch.sum(sims_pos_sigm, dim=-1) + 1
 
-        ap = torch.zeros(1).to(embeddings.device)
         g = batch_size // num_classes_batch
+        ap = torch.zeros(batch_size).to(embeddings.device)
         for i in range(num_classes_batch):
-            pos_divide = torch.sum(
-                sims_pos_ranks[i] / sims_ranks[i * g : (i + 1) * g, i * g : (i + 1) * g]
-            )
-            ap = ap + (pos_divide / g) / batch_size
+            for j in range(g):
+                pos_rank = sims_pos_ranks[i, j]
+                all_rank = sims_ranks[i * g + j, i * g: (i + 1) * g]
+                ap[i * g + j] = torch.sum(pos_rank / all_rank) / g
 
         loss = 1 - ap
+
         return {
-            "loss": {
+            "ap_loss": {
                 "losses": loss,
-                "indices": None,
-                "reduction_type": "already_reduced",
+                "indices": c_f.torch_arange_from_size(loss),
+                "reduction_type": "element",
             }
-        }
+        }

tests/losses/test_smooth_ap_loss.py

@@ -187,7 +187,7 @@ def test_smooth_ap_loss(self):
                     labels.extend([i for _ in range(HYPERPARAMETERS["num_id"])])
 
                 labels = torch.tensor(labels)
-                output2 = loss2.compute_loss(
+                output2 = loss2.forward(
                     rand_tensor, labels, None, rand_tensor, labels
-                )["loss"]["losses"]
+                )
                 self.assertTrue(torch.isclose(output, output2))