update

top5_error_rate.py

@@ -44,39 +44,39 @@ class Top5ErrorRate(evaluate.Metric):
             inputs_description=_KWARGS_DESCRIPTION,
             features=datasets.Features(
                 {
-                    "predictions": datasets.Sequence(datasets.Sequence(datasets.Value("float32"))),
-                    "references": datasets.Sequence(datasets.Value("int32")),
+                    'predictions': datasets.Sequence(datasets.Sequence(datasets.Value('float32'))),
+                    'references': datasets.Sequence(datasets.Value('int32')),
                 }
-                if self.config_name == "multilabel"
+                if self.config_name == 'multilabel'
                 else {
-                    "predictions": datasets.Sequence(datasets.Value("float32")),
-                    "references": datasets.Value("int32"),
+                    'predictions': datasets.Sequence(datasets.Value('float32')),
+                    'references': datasets.Value('int32'),
                 }
             ),
             reference_urls=[],
         )
 
     def _compute(
-            self,
-            *,
-            predictions: list[list[float]] = None,
-            references: list[int] = None,
-            **kwargs,
+        self,
+        *,
+        predictions: list[list[float]] = None,
+        references: list[int] = None,
+        **kwargs,
     ) -> Dict[str, Any]:
         # to numpy array
-        outputs = np.array(predictions, dtype=np.float32)
+        outputs = np.array(predictions)
         labels = np.array(references)
 
         # Top-1 ACC
         pred = outputs.argmax(axis=1)
         acc = (pred == labels).mean()
 
-        # Top-5 Error Rate
-        top5_indices = outputs.argsort(axis=1)[:, -5:]
-        correct = (labels.reshape(-1, 1) == top5_indices).any(axis=1)
+        # Top-5 Error rate
+        top5_indices = np.argpartition(outputs, -5, axis=1)[:, -5:]
+
+        # 使用广播机制直接比较
+        # 使用np.any的axis参数直接在最后一个维度上检查是否存在匹配
+        correct = np.any(top5_indices == labels[:, np.newaxis], axis=1)
         top5_error_rate = 1 - correct.mean()
 
-        return {
-            "accuracy": float(acc),
-            "top5_error_rate": float(top5_error_rate)
-        }
+        return {'accuracy': float(acc), 'top5_error_rate': float(top5_error_rate)}