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())
class Tester:
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 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 __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 __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 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 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 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()
class DefaultTrainer(SimpleTrainer):
"""
A trainer with default training logic. Compared to `SimpleTrainer`, it
contains the following logic in addition:
1. Create model, optimizer, scheduler, dataloader from the given config.
2. Load a checkpoint or `cfg.MODEL.WEIGHTS`, if exists.
3. Register a few common hooks.
It is created to simplify the **standard model training workflow** and
reduce code boilerplate for users who only need the standard training
workflow, with standard features. It means this class makes *many
assumptions* about your training logic that may easily become invalid in
a new research. In fact, any assumptions beyond those made in the
:class:`SimpleTrainer` are too much for research.
The code of this class has been annotated about restrictive assumptions
it mades. When they do not work for you, you're encouraged to:
1. Overwrite methods of this class, OR:
2. Use :class:`SimpleTrainer`, which only does minimal SGD training and
nothing else. You can then add your own hooks if needed. OR:
3. Write your own training loop similar to `tools/plain_train_net.py`.
Also note that the behavior of this class, like other functions/classes in
this file, is not stable, since it is meant to represent the "common
default behavior".
It is only guaranteed to work well with the standard models and training
workflow in FsDet.
To obtain more stable behavior, write your own training logic with other
public APIs.
Attributes:
scheduler:
checkpointer (DetectionCheckpointer):
cfg (CfgNode):
Examples:
.. code-block:: python
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
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 resume_or_load(self, resume=True):
"""
If `resume==True`, and last checkpoint exists, resume from it.
Otherwise, load a model specified by the config.
Args:
resume (bool): whether to do resume or not
"""
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
self.start_iter = (self.checkpointer.resume_or_load(
self.cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1)
def build_hooks(self):
"""
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
Returns:
list[HookBase]:
"""
cfg = self.cfg.clone()
cfg.defrost()
cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
ret = [
hooks.IterationTimer(),
hooks.LRScheduler(self.optimizer, self.scheduler),
hooks.PreciseBN(
# Run at the same freq as (but before) evaluation.
cfg.TEST.EVAL_PERIOD,
self.model,
# Build a new data loader to not affect training
self.build_train_loader(cfg),
cfg.TEST.PRECISE_BN.NUM_ITER,
) if cfg.TEST.PRECISE_BN.ENABLED and get_bn_modules(self.model)
else None,
]
# Do PreciseBN before checkpointer, because it updates the model and need to
# be saved by checkpointer.
# This is not always the best: if checkpointing has a different frequency,
# some checkpoints may have more precise statistics than others.
if comm.is_main_process():
ret.append(
hooks.PeriodicCheckpointer(self.checkpointer,
cfg.SOLVER.CHECKPOINT_PERIOD))
def test_and_save_results():
self._last_eval_results = self.test(self.cfg, self.model)
return self._last_eval_results
# Do evaluation after checkpointer, because then if it fails,
# we can use the saved checkpoint to debug.
ret.append(hooks.EvalHook(cfg.TEST.EVAL_PERIOD, test_and_save_results))
if comm.is_main_process():
# run writers in the end, so that evaluation metrics are written
ret.append(hooks.PeriodicWriter(self.build_writers()))
return ret
def build_writers(self):
"""
Build a list of writers to be used. By default it contains
writers that write metrics to the screen,
a json file, and a tensorboard event file respectively.
If you'd like a different list of writers, you can overwrite it in
your trainer.
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
It is now implemented by:
.. code-block:: python
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Assume the default print/log frequency.
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
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
@classmethod
def build_model(cls, cfg):
"""
Returns:
torch.nn.Module:
It now calls :func:`fsdet.modeling.build_model`.
Overwrite it if you'd like a different model.
"""
model = build_model(cfg)
logger = logging.getLogger(__name__)
if not cfg.MUTE_HEADER:
logger.info("Model:\n{}".format(model))
return model
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`fsdet.solver.build_optimizer`.
Overwrite it if you'd like a different optimizer.
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer):
"""
It now calls :func:`fsdet.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`fsdet.data.build_detection_train_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_detection_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`fsdet.data.build_detection_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_detection_test_loader(cfg, dataset_name)
@classmethod
def build_evaluator(cls, cfg, dataset_name):
"""
Returns:
DatasetEvaluator
It is not implemented by default.
"""
raise NotImplementedError(
"Please either implement `build_evaluator()` in subclasses, or pass "
"your evaluator as arguments to `DefaultTrainer.test()`.")
@classmethod
def test(cls, cfg, model, evaluators=None):
"""
Args:
cfg (CfgNode):
model (nn.Module):
evaluators (list[DatasetEvaluator] or None): if None, will call
:meth:`build_evaluator`. Otherwise, must have the same length as
`cfg.DATASETS.TEST`.
Returns:
dict: a dict of result metrics
"""
logger = logging.getLogger(__name__)
if isinstance(evaluators, DatasetEvaluator):
evaluators = [evaluators]
if evaluators is not None:
assert len(
cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
len(cfg.DATASETS.TEST), len(evaluators))
results = OrderedDict()
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
data_loader = cls.build_test_loader(cfg, dataset_name)
# When evaluators are passed in as arguments,
# implicitly assume that evaluators can be created before data_loader.
if evaluators is not None:
evaluator = evaluators[idx]
else:
try:
evaluator = cls.build_evaluator(cfg, dataset_name)
except NotImplementedError:
logger.warn(
"No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
"or implement its `build_evaluator` method.")
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model, data_loader, evaluator)
results[dataset_name] = results_i
if comm.is_main_process():
assert isinstance(
results_i, dict
), "Evaluator must return a dict on the main process. Got {} instead.".format(
results_i)
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
print_csv_format(results_i)
if len(results) == 1:
results = list(results.values())[0]
return results