class Trainer(DefaultTrainer):
"""
This is the same Trainer except that we rewrite the
`build_train_loader`/`resume_or_load` method.
"""
def build_hooks(self):
"""
Replace `DetectionCheckpointer` with `AdetCheckpointer`.
Build a list of default hooks, including timing, evaluation,
checkpointing, lr scheduling, precise BN, writing events.
"""
ret = super().build_hooks()
for i in range(len(ret)):
if isinstance(ret[i], hooks.PeriodicCheckpointer):
self.checkpointer = AdetCheckpointer(
self.model,
self.cfg.OUTPUT_DIR,
optimizer=self.optimizer,
scheduler=self.scheduler,
)
ret[i] = hooks.PeriodicCheckpointer(self.checkpointer, self.cfg.SOLVER.CHECKPOINT_PERIOD)
return ret
def resume_or_load(self, resume=True):
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS, resume=resume)
if resume and self.checkpointer.has_checkpoint():
self.start_iter = checkpoint.get("iteration", -1) + 1
def train_loop(self, start_iter: int, max_iter: int):
"""
Args:
start_iter, max_iter (int): See docs above
"""
logger = logging.getLogger("adet.trainer")
logger.info("Starting training from iteration {}".format(start_iter))
self.iter = self.start_iter = start_iter
self.max_iter = max_iter
with EventStorage(start_iter) as self.storage:
self.before_train()
for self.iter in range(start_iter, max_iter):
self.before_step()
self.run_step()
self.after_step()
self.after_train()
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
self.train_loop(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_train_loader(cls, cfg):
"""
Returns:
iterable
It calls :func:`detectron2.data.build_detection_train_loader` with a customized
DatasetMapper, which adds categorical labels as a semantic mask.
"""
mapper = DatasetMapperWithBasis(cfg, True)
return build_detection_train_loader(cfg, mapper=mapper)
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
"""
Create evaluator(s) for a given dataset.
This uses the special metadata "evaluator_type" associated with each builtin dataset.
For your own dataset, you can simply create an evaluator manually in your
script and do not have to worry about the hacky if-else logic here.
"""
if output_folder is None:
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
evaluator_list = []
evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
evaluator_list.append(
SemSegEvaluator(
dataset_name,
distributed=True,
num_classes=cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
ignore_label=cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
output_dir=output_folder,
)
)
if evaluator_type in ["coco", "coco_panoptic_seg"]:
evaluator_list.append(COCOEvaluator(dataset_name, cfg, True, output_folder))
if evaluator_type == "coco_panoptic_seg":
evaluator_list.append(COCOPanopticEvaluator(dataset_name, output_folder))
if evaluator_type == "pascal_voc":
return PascalVOCDetectionEvaluator(dataset_name)
if evaluator_type == "lvis":
return LVISEvaluator(dataset_name, cfg, True, output_folder)
if evaluator_type == "text":
return TextEvaluator(dataset_name, cfg, True, output_folder)
if len(evaluator_list) == 0:
raise NotImplementedError(
"no Evaluator for the dataset {} with the type {}".format(
dataset_name, evaluator_type
)
)
if len(evaluator_list) == 1:
return evaluator_list[0]
return DatasetEvaluators(evaluator_list)
@classmethod
def test_with_TTA(cls, cfg, model):
logger = logging.getLogger("adet.trainer")
# In the end of training, run an evaluation with TTA
# Only support some R-CNN models.
logger.info("Running inference with test-time augmentation ...")
model = GeneralizedRCNNWithTTA(cfg, model)
evaluators = [
cls.build_evaluator(
cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
)
for name in cfg.DATASETS.TEST
]
res = cls.test(cfg, model, evaluators)
res = OrderedDict({k + "_TTA": v for k, v in res.items()})
return res
class MyTrainer(DefaultTrainer):
"""
This is the same Trainer except that we rewrite the
`build_train_loader` method.
"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
Use the custom checkpointer, which loads other backbone models
with matching heuristics.
"""
# Assume these objects must be constructed in this order.
dprint("build model")
model = self.build_model(cfg)
dprint('build optimizer')
optimizer = self.build_optimizer(cfg, model)
dprint("build train loader")
data_loader = self.build_train_loader(cfg)
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
if isinstance(data_loader, AspectRatioGroupedDataset):
dataset_len = len(data_loader.dataset.dataset)
iters_per_epoch = dataset_len // images_per_batch
else:
dataset_len = len(data_loader.dataset)
iters_per_epoch = dataset_len // images_per_batch
self.iters_per_epoch = iters_per_epoch
total_iters = cfg.SOLVER.TOTAL_EPOCHS * iters_per_epoch
dprint("images_per_batch: ", images_per_batch)
dprint("dataset length: ", dataset_len)
dprint("iters per epoch: ", iters_per_epoch)
dprint("total iters: ", total_iters)
# 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)
super(DefaultTrainer, self).__init__(model, data_loader, optimizer)
self.scheduler = self.build_lr_scheduler(cfg,
optimizer,
total_iters=total_iters)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = AdetCheckpointer(
# 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 = total_iters # NOTE: ignore 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, load all checkpointables
(eg. optimizer and scheduler) and update iteration counter.
Otherwise, load the model specified by the config (skip all checkpointables) and start from
the first iteration.
Args:
resume (bool): whether to do resume or not
"""
checkpoint = self.checkpointer.resume_or_load(self.cfg.MODEL.WEIGHTS,
resume=resume)
if resume and self.checkpointer.has_checkpoint():
self.start_iter = checkpoint.get("iteration", -1) + 1
# The checkpoint stores the training iteration that just finished, thus we start
# at the next iteration (or iter zero if there's no checkpoint).
@classmethod
def build_optimizer(cls, cfg, model):
return solver_utils.my_build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, total_iters):
# NOTE: here we ignore MAX_ITER and STEPS in config, and use
# TOTAL_EPOCHS * len(data_loader), REL_STEPS instead
return solver_utils.build_lr_scheduler(cfg,
optimizer,
total_iters=total_iters)
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():
if cfg.SOLVER.CHECKPOINT_BY_EPOCH:
ckpt_period = cfg.SOLVER.CHECKPOINT_PERIOD * self.iters_per_epoch
else:
ckpt_period = cfg.SOLVER.CHECKPOINT_PERIOD
ret.append(
MyPeriodicCheckpointer(self.checkpointer,
ckpt_period,
max_to_keep=cfg.SOLVER.get(
"NUM_CKPT_KEEP", 5),
iters_per_epoch=self.iters_per_epoch))
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(),
period=cfg.TRAIN.get("PRINT_FREQ", 100)))
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:
::
return [
CommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(self.cfg.OUTPUT_DIR),
]
"""
# Here the default print/log frequency of each writer is used.
tb_logdir = osp.join(self.cfg.OUTPUT_DIR, "tb")
mmcv.mkdir_or_exist(tb_logdir)
cfg = self.cfg
if not self.cfg.get("RESUME", False):
old_tb_logdir = osp.join(cfg.OUTPUT_DIR, "tb_old")
mmcv.mkdir_or_exist(old_tb_logdir)
os.system("mv -v {} {}".format(osp.join(tb_logdir, "events.*"),
old_tb_logdir))
tbx_event_writer = MyTensorboardXWriter(tb_logdir,
backend="tensorboardX")
self.tbx_writer = tbx_event_writer._writer # NOTE: we want to write some non-scalar data
return [
# It may not always print what you want to see, since it prints "common" metrics only.
MyCommonMetricPrinter(self.max_iter),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json")),
tbx_event_writer,
]
def train_loop(self, start_iter: int, max_iter: int):
"""
Args:
start_iter, max_iter (int): See docs above
"""
logger = logging.getLogger("adet.trainer")
logger.info("Starting training from iteration {}".format(start_iter))
self.iter = self.start_iter = start_iter
self.max_iter = max_iter
with EventStorage(start_iter) as self.storage:
self.before_train()
for self.iter in range(start_iter, max_iter):
self.before_step()
self.run_step()
self.after_step()
self.after_train()
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
self.train_loop(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_train_loader(cls, cfg):
"""
Returns:
iterable
It calls :func:`detectron2.data.build_detection_train_loader` with a customized
DatasetMapper, which adds categorical labels as a semantic mask.
"""
mapper = MyDatasetMapperWithBasis(cfg, True)
return build_detection_train_loader(cfg, mapper)
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
"""
Create evaluator(s) for a given dataset.
This uses the special metadata "evaluator_type" associated with each builtin dataset.
For your own dataset, you can simply create an evaluator manually in your
script and do not have to worry about the hacky if-else logic here.
"""
if output_folder is None:
model_name = osp.basename(cfg.MODEL.WEIGHTS).split(".")[0]
output_folder = osp.join(
cfg.OUTPUT_DIR,
"inference_{}/{}".format(model_name, dataset_name))
evaluator_list = []
evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
if evaluator_type in ["sem_seg", "coco_panoptic_seg"]:
evaluator_list.append(
SemSegEvaluator(
dataset_name,
distributed=True,
num_classes=cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES,
ignore_label=cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE,
output_dir=output_folder,
))
if evaluator_type in ["coco", "coco_panoptic_seg", "coco_bop"]:
evaluator_list.append(
MyCOCOEvaluator(dataset_name, cfg, True, output_folder))
# evaluator_list.append(COCOEvaluator(dataset_name, cfg, True, output_folder))
if evaluator_type == "coco_panoptic_seg":
evaluator_list.append(
COCOPanopticEvaluator(dataset_name, output_folder))
if evaluator_type == "pascal_voc":
return PascalVOCDetectionEvaluator(dataset_name)
if evaluator_type == "lvis":
return LVISEvaluator(dataset_name, cfg, True, output_folder)
if evaluator_type == "text":
return TextEvaluator(dataset_name, cfg, True, output_folder)
if len(evaluator_list) == 0:
raise NotImplementedError(
"no Evaluator for the dataset {} with the type {}".format(
dataset_name, evaluator_type))
if len(evaluator_list) == 1:
return evaluator_list[0]
return DatasetEvaluators(evaluator_list)
@classmethod
def test_with_TTA(cls, cfg, model):
logger = logging.getLogger("adet.trainer")
# In the end of training, run an evaluation with TTA
# Only support some R-CNN models.
logger.info("Running inference with test-time augmentation ...")
model = GeneralizedRCNNWithTTA(cfg, model)
evaluators = [
cls.build_evaluator(cfg,
name,
output_folder=os.path.join(
cfg.OUTPUT_DIR, "inference_TTA"))
for name in cfg.DATASETS.TEST
]
res = cls.test(cfg, model, evaluators)
res = OrderedDict({k + "_TTA": v for k, v in res.items()})
return res
@classmethod
def build_test_loader(cls, cfg, dataset_name):
"""
Returns:
iterable
It now calls :func:`detectron2.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 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