detectron2 打印各类ap50 & 保存map50最佳模型 & 训练时验证

打印各类ap50，参考下面链接修改detectron2/evaluation/coco_evaluation.py 中的代码https://github.com/facebookresearch/detectron2/issues/5101#issuecomment-1843388809

diff --git a/detectron2/evaluation/coco_evaluation.py b/detectron2/evaluation/coco_evaluation.py
index fe8142c..5b6db0a 100644
--- a/detectron2/evaluation/coco_evaluation.py
+++ b/detectron2/evaluation/coco_evaluation.py
@@ -362,12 +362,13 @@ class COCOEvaluator(DatasetEvaluator):
         precisions = coco_eval.eval["precision"]
         # precision has dims (iou, recall, cls, area range, max dets)
         assert len(class_names) == precisions.shape[2]
+        assert coco_eval.params.iouThrs[0] == 0.5
 
         results_per_category = []
         for idx, name in enumerate(class_names):
             # area range index 0: all area ranges
             # max dets index -1: typically 100 per image
-            precision = precisions[:, :, idx, 0, -1]
+            precision = precisions[0, :, idx, 0, -1]
             precision = precision[precision > -1]
             ap = np.mean(precision) if precision.size else float("nan")
             results_per_category.append(("{}".format(name), float(ap * 100)))
@@ -380,12 +381,12 @@ class COCOEvaluator(DatasetEvaluator):
             results_2d,
             tablefmt="pipe",
             floatfmt=".3f",
-            headers=["category", "AP"] * (N_COLS // 2),
+            headers=["category", "AP50"] * (N_COLS // 2),
             numalign="left",
         )
-        self._logger.info("Per-category {} AP: \n".format(iou_type) + table)
+        self._logger.info("Per-category {} AP50: \n".format(iou_type) + table)
 
-        results.update({"AP-" + name: ap for name, ap in results_per_category})
+        results.update({"AP50-" + name: ap for name, ap in results_per_category})
         return results

在 train.py 的main函数中添加代码trainer.register_hooks，注册BestCheckpointer

trainer = Trainer(cfg)
# here to add
trainer.register_hooks([hooks.BestCheckpointer(eval_period=cfg.TEST.EVAL_PERIOD,checkpointer=trainer.checkpointer,val_metric= "bbox/AP50", mode="max")])
trainer.resume_or_load(resume=args.resume)
return trainer.train()

demo.yaml 文件中修改 EVAL_PERIOD，调整训练过程何时验证，这里我的数据集计算出来是160个 iteration 为一个epoch，故设为160

TEST:
  EVAL_PERIOD: 160