def test(self, ckpt):
self.check_pointer._load_model(self.check_pointer._load_file(ckpt))
print('evaluating checkpoint {}'.format(ckpt))
res = Trainer.test(self.cfg, self.model)
if comm.is_main_process():
verify_results(self.cfg, res)
print(res)
key = 'nAP50'
if key not in res['bbox']:
key = 'AP'
if (self.best_res is None) or (
self.best_res is not None and self.best_res['bbox'][key] <
res['bbox'][key]): # best result by the above key
self.best_res = res
self.best_file = ckpt
print('best results from checkpoint {}'.format(self.best_file))
print(self.best_res)
self.all_res["best_file"] = self.best_file
self.all_res["best_res"] = self.best_res
self.all_res[ckpt] = res
os.makedirs(os.path.join(self.cfg.OUTPUT_DIR, 'inference'),
exist_ok=True)
with open(
os.path.join(self.cfg.OUTPUT_DIR, 'inference',
'all_res.json'), 'w') as fp:
json.dump(self.all_res, fp)
def test(self, ckpt):
self.check_pointer._load_model(self.check_pointer._load_file(ckpt))
print("evaluating checkpoint {}".format(ckpt))
res = Trainer.test(self.cfg, self.model)
if comm.is_main_process():
verify_results(self.cfg, res)
print(res)
if (self.best_res is None) or (
self.best_res is not None
and self.best_res["bbox"]["AP"] < res["bbox"]["AP"]
):
self.best_res = res
self.best_file = ckpt
print("best results from checkpoint {}".format(self.best_file))
print(self.best_res)
self.all_res["best_file"] = self.best_file
self.all_res["best_res"] = self.best_res
self.all_res[ckpt] = res
os.makedirs(
os.path.join(self.cfg.OUTPUT_DIR, "inference"), exist_ok=True
)
with open(
os.path.join(self.cfg.OUTPUT_DIR, "inference", "all_res.json"),
"w",
) as fp:
json.dump(self.all_res, fp)
def main(args):
cfg = setup(args)
if args.eval_only:
model = Trainer.build_model(cfg)
if args.eval_iter != -1:
# load checkpoint at specified iteration
ckpt_file = os.path.join(
cfg.OUTPUT_DIR, "model_{:07d}.pth".format(args.eval_iter - 1))
resume = False
else:
# load checkpoint at last iteration
ckpt_file = cfg.MODEL.WEIGHTS
resume = True
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
ckpt_file, resume=resume)
res = Trainer.test(cfg, model)
if comm.is_main_process():
verify_results(cfg, res)
# save evaluation results in json
os.makedirs(os.path.join(cfg.OUTPUT_DIR, "inference"),
exist_ok=True)
with open(
os.path.join(cfg.OUTPUT_DIR, "inference",
"res_final.json"),
"w",
) as fp:
json.dump(res, fp)
return res
elif args.eval_all:
tester = Tester(cfg)
all_ckpts = sorted(tester.check_pointer.get_all_checkpoint_files())
for i, ckpt in enumerate(all_ckpts):
ckpt_iter = ckpt.split("model_")[-1].split(".pth")[0]
if ckpt_iter.isnumeric() and int(ckpt_iter) + 1 < args.start_iter:
# skip evaluation of checkpoints before start iteration
continue
if args.end_iter != -1:
if (not ckpt_iter.isnumeric()
or int(ckpt_iter) + 1 > args.end_iter):
# skip evaluation of checkpoints after end iteration
break
tester.test(ckpt)
return best_res
elif args.eval_during_train:
tester = Tester(cfg)
saved_checkpoint = None
while True:
if tester.check_pointer.has_checkpoint():
current_ckpt = tester.check_pointer.get_checkpoint_file()
if (saved_checkpoint is None
or current_ckpt != saved_checkpoint):
saved_checkpoint = current_ckpt
tester.test(current_ckpt)
time.sleep(10)
else:
if comm.is_main_process():
print(
"Please specify --eval-only, --eval-all, or --eval-during-train"
)
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
super().train(self.start_iter, self.max_iter)
if hasattr(self, "_last_eval_results") and comm.is_main_process():
verify_results(self.cfg, self._last_eval_results)
return self._last_eval_results
def main(args):
cfg = setup(args)
if args.eval_only:
model = Trainer.build_model(cfg)
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
res = Trainer.test(cfg, model)
if comm.is_main_process():
verify_results(cfg, res)
return res
"""
If you'd like to do anything fancier than the standard training logic,
consider writing your own training loop or subclassing the trainer.
"""
trainer = Trainer(cfg)
trainer.resume_or_load(resume=args.resume)
return trainer.train()