update segment training (#57)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: ayush chaurasia <ayush.chaurarsia@gmail.com>
2022-11-29 05:30:08 -06:00 · 2022-11-29 05:30:08 -06:00 · 3a241e4cea
commit 3a241e4cea
parent d0b0fe2592
14 changed files with 460 additions and 144 deletions
--- a/ultralytics/yolo/v8/classify/val.py
+++ b/ultralytics/yolo/v8/classify/val.py
@ -23,3 +23,7 @@ class ClassificationValidator(BaseValidator):
        acc = torch.stack((self.correct[:, 0], self.correct.max(1).values), dim=1)  # (top1, top5) accuracy
        top1, top5 = acc.mean(0).tolist()
        return {"top1": top1, "top5": top5, "fitness": top5}
+
+    @property
+    def metric_keys(self):
+        return ["top1", "top5"]