def get(config_path, trained: bool = False):
"""
Get a model specified by relative path under FsDet's official ``configs/`` directory.
Args:
config_path (str): config file name relative to FsDet's "configs/"
directory, e.g., "COCO-detection/faster_rcnn_R_101_FPN_ft_all_1shot.yaml"
trained (bool): If True, will initialize the model with the trained model zoo weights.
If False, the checkpoint specified in the config file's ``MODEL.WEIGHTS`` is used
instead; this will typically (though not always) initialize a subset of weights using
an ImageNet pre-trained model, while randomly initializing the other weights.
Example:
.. code-block:: python
from fsdet import model_zoo
model = model_zoo.get("COCO-detection/faster_rcnn_R_101_FPN_ft_all_1shot.yaml", trained=True)
"""
cfg_file = get_config_file(config_path)
cfg = get_cfg()
cfg.merge_from_file(cfg_file)
if trained:
cfg.MODEL.WEIGHTS = get_checkpoint_url(config_path)
if not torch.cuda.is_available():
cfg.MODEL.DEVICE = "cpu"
model = build_model(cfg)
DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS)
return model
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
"""
# Assume these objects must be constructed in this order.
model = self.build_model(cfg)
optimizer = self.build_optimizer(cfg, model)
data_loader = self.build_train_loader(cfg)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
super().__init__(model, data_loader, optimizer)
self.scheduler = self.build_lr_scheduler(cfg, optimizer)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = DetectionCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
self.max_iter = cfg.SOLVER.MAX_ITER
self.cfg = cfg
self.register_hooks(self.build_hooks())
def __init__(self, cfg):
self.cfg = cfg
self.model = Trainer.build_model(cfg)
self.check_pointer = DetectionCheckpointer(self.model,
save_dir=cfg.OUTPUT_DIR)
self.best_res = None
self.best_file = None
self.all_res = {}
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 __init__(self, cfg):
self.cfg = cfg.clone() # cfg can be modified by model
self.model = build_model(self.cfg)
self.model.eval()
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
checkpointer = DetectionCheckpointer(self.model)
checkpointer.load(cfg.MODEL.WEIGHTS)
self.transform_gen = T.ResizeShortestEdge(
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST],
cfg.INPUT.MAX_SIZE_TEST)
self.input_format = cfg.INPUT.FORMAT
assert self.input_format in ["RGB", "BGR"], self.input_format
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()