def stage_main(args, cfg, build):
cfg.merge_from_list(args.opts)
cfg, logger = default_setup(cfg, args)
model_build_func = build
"""
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, model_build_func)
trainer.resume_or_load(resume=args.resume)
if args.eval_only:
DefaultCheckpointer(
trainer.model, save_dir=cfg.OUTPUT_DIR, resume=args.resume).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
res = Trainer.test(cfg, trainer.model)
if comm.is_main_process():
verify_results(cfg, res)
if cfg.TEST.AUG.ENABLED:
res.update(Trainer.test_with_TTA(cfg, trainer.model))
return res
# check wheather worksapce has enough storeage space
# assume that a single dumped model is 700Mb
file_sys = os.statvfs(cfg.OUTPUT_DIR)
free_space_Gb = (file_sys.f_bfree * file_sys.f_frsize) / 2**30
eval_space_Gb = (cfg.SOLVER.LR_SCHEDULER.MAX_ITER // cfg.SOLVER.CHECKPOINT_PERIOD) * 700 / 2**10
if eval_space_Gb > free_space_Gb:
logger.warning(f"{Fore.RED}Remaining space({free_space_Gb}GB) "
f"is less than ({eval_space_Gb}GB){Style.RESET_ALL}")
if cfg.TEST.AUG.ENABLED:
trainer.register_hooks(
[hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
)
trainer.train()
if comm.is_main_process() and cfg.MODEL.AS_PRETRAIN:
# convert last ckpt to pretrain format
convert_to_pretrained_model(
input=os.path.join(cfg.OUTPUT_DIR, "model_final.pth"),
save_path=os.path.join(cfg.OUTPUT_DIR, "model_final_pretrain_weight.pkl")
)
def main(args):
config.merge_from_list(args.opts)
cfg = setup(args)
model = build_model(cfg)
logger.info("Model:\n{}".format(model))
if args.eval_only:
DefaultCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume)
return do_test(cfg, model)
distributed = comm.get_world_size() > 1
if distributed:
model = DistributedDataParallel(model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False)
do_train(cfg, model)
return do_test(cfg, model)
def __init__(self, cfg, model_build_func):
"""
Args:
cfg (BaseConfig):
"""
logger = logging.getLogger("cvpods")
if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2
setup_logger()
self.start_iter = 0
data_loader = self.build_train_loader(cfg)
maybe_adjust_epoch_and_iter(cfg, data_loader)
self.max_iter = cfg.SOLVER.LR_SCHEDULER.MAX_ITER
self.max_epoch = cfg.SOLVER.LR_SCHEDULER.MAX_EPOCH
model = model_build_func(cfg)
model = maybe_convert_module(model)
logger.info(f"Model structure: {model}")
# Assume these objects must be constructed in this order.
optimizer = self.build_optimizer(cfg, model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
if cfg.TRAINER.FP16.ENABLED:
if cfg.TRAINER.FP16.TYPE == "APEX":
model, optimizer = amp.initialize(
model, optimizer, opt_level=cfg.TRAINER.FP16.OPTS.OPT_LEVEL
)
model = DistributedDataParallel(
model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True
)
# TODO: @wangfeng02, `batch_subdivisions`
super().__init__(model, data_loader, optimizer, cfg.SOLVER.BATCH_SUBDIVISIONS)
if not cfg.SOLVER.LR_SCHEDULER.get("EPOCH_WISE", False):
epoch_iters = -1
else:
epoch_iters = cfg.SOLVER.LR_SCHEDULER.get("EPOCH_ITERS")
logger.warning(f"Setup LR Scheduler in EPOCH mode: {epoch_iters}")
self.scheduler = self.build_lr_scheduler(
cfg, optimizer, epoch_iters=epoch_iters)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
optional = {}
if cfg.TRAINER.FP16.ENABLED:
optional["amp"] = amp
self.checkpointer = DefaultCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler,
**optional,
)
self.cfg = cfg
self.register_hooks(self.build_hooks())
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 cvpods.
To obtain more stable behavior, write your own training logic with other public APIs.
Attributes:
scheduler:
checkpointer (DefaultCheckpointer):
cfg (BaseConfig):
Examples:
.. code-block:: python
trainer = DefaultTrainer(cfg)
trainer.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
trainer.train()
"""
def __init__(self, cfg, model_build_func):
"""
Args:
cfg (BaseConfig):
"""
logger = logging.getLogger("cvpods")
if not logger.isEnabledFor(logging.INFO): # setup_logger is not called for d2
setup_logger()
self.start_iter = 0
data_loader = self.build_train_loader(cfg)
maybe_adjust_epoch_and_iter(cfg, data_loader)
self.max_iter = cfg.SOLVER.LR_SCHEDULER.MAX_ITER
self.max_epoch = cfg.SOLVER.LR_SCHEDULER.MAX_EPOCH
model = model_build_func(cfg)
model = maybe_convert_module(model)
logger.info(f"Model structure: {model}")
# Assume these objects must be constructed in this order.
optimizer = self.build_optimizer(cfg, model)
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
if cfg.TRAINER.FP16.ENABLED:
if cfg.TRAINER.FP16.TYPE == "APEX":
model, optimizer = amp.initialize(
model, optimizer, opt_level=cfg.TRAINER.FP16.OPTS.OPT_LEVEL
)
model = DistributedDataParallel(
model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True
)
# TODO: @wangfeng02, `batch_subdivisions`
super().__init__(model, data_loader, optimizer, cfg.SOLVER.BATCH_SUBDIVISIONS)
if not cfg.SOLVER.LR_SCHEDULER.get("EPOCH_WISE", False):
epoch_iters = -1
else:
epoch_iters = cfg.SOLVER.LR_SCHEDULER.get("EPOCH_ITERS")
logger.warning(f"Setup LR Scheduler in EPOCH mode: {epoch_iters}")
self.scheduler = self.build_lr_scheduler(
cfg, optimizer, epoch_iters=epoch_iters)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
optional = {}
if cfg.TRAINER.FP16.ENABLED:
optional["amp"] = amp
self.checkpointer = DefaultCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
model,
cfg.OUTPUT_DIR,
optimizer=optimizer,
scheduler=self.scheduler,
**optional,
)
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
"""
self.checkpointer.resume = resume
# 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
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(), period=self.cfg.GLOBAL.LOG_INTERVAL
))
# Put `PeriodicDumpLog` after writers so that can dump all the files,
# including the files generated by 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, self.max_epoch)
if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
verify_results(self.cfg, self._last_eval_results)
return self._last_eval_results
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
"""
return build_optimizer(cfg, model)
@classmethod
def build_lr_scheduler(cls, cfg, optimizer, **kwargs):
"""
It now calls :func:`cvpods.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, **kwargs)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`cvpods.data.build_train_loader`.
Overwrite it if you'd like a different data loader.:w
"""
return build_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`cvpods.data.build_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_test_loader(cfg)
@classmethod
def build_evaluator(cls, cfg, dataset_name):
"""
Returns:
DatasetEvaluator or None
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, output_folder=None):
"""
Args:
cfg (BaseConfig):
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)
# 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, data_loader.dataset, output_folder=output_folder)
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
def do_train(cfg, model, resume=False):
model.train()
optimizer = build_optimizer(cfg, model)
scheduler = build_lr_scheduler(cfg, optimizer)
checkpointer = DefaultCheckpointer(
model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=scheduler
)
start_iter = (
checkpointer.resume_or_load(cfg.MODEL.WEIGHTS, resume=resume).get("iteration", -1) + 1
)
max_iter = cfg.SOLVER.MAX_ITER
periodic_checkpointer = PeriodicCheckpointer(
checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter
)
writers = (
[
CommonMetricPrinter(max_iter),
JSONWriter(os.path.join(cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(cfg.OUTPUT_DIR),
]
if comm.is_main_process()
else []
)
# compared to "train_net.py", we do not support accurate timing and
# precise BN here, because they are not trivial to implement
data_loader = build_train_loader(cfg)
logger.info("Starting training from iteration {}".format(start_iter))
with EventStorage(start_iter) as storage:
for data, iteration in zip(data_loader, range(start_iter, max_iter)):
iteration = iteration + 1
storage.step()
loss_dict = model(data)
losses = sum(loss for loss in loss_dict.values())
assert torch.isfinite(losses).all(), loss_dict
loss_dict_reduced = {k: v.item() for k, v in comm.reduce_dict(loss_dict).items()}
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
if comm.is_main_process():
storage.put_scalars(total_loss=losses_reduced, **loss_dict_reduced)
optimizer.zero_grad()
losses.backward()
optimizer.step()
storage.put_scalar("lr", optimizer.param_groups[0]["lr"], smoothing_hint=False)
scheduler.step()
if (
cfg.TEST.EVAL_PERIOD > 0
and iteration % cfg.TEST.EVAL_PERIOD == 0
and iteration != max_iter
):
do_test(cfg, model)
# Compared to "train_net.py", the test results are not dumped to EventStorage
comm.synchronize()
if iteration - start_iter > 5 and (iteration % 20 == 0 or iteration == max_iter):
for writer in writers:
writer.write()
periodic_checkpointer.step(iteration)
help="output directory for the converted caffe2 model")
parser.add_argument(
"opts",
help="Modify config options using the command-line",
default=None,
nargs=argparse.REMAINDER,
)
args = parser.parse_args()
logger = setup_logger()
logger.info("Command line arguments: " + str(args))
cfg = setup_cfg(args)
# create a torch model
torch_model = build_model(cfg)
DefaultCheckpointer(torch_model).resume_or_load(cfg.MODEL.WEIGHTS)
# get a sample data
data_loader = build_test_loader(cfg, cfg.DATASETS.TEST[0])
first_batch = next(iter(data_loader))
# convert and save caffe2 model
caffe2_model = export_caffe2_model(cfg, torch_model, first_batch)
caffe2_model.save_protobuf(args.output)
# draw the caffe2 graph
caffe2_model.save_graph(os.path.join(args.output, "model.svg"),
inputs=first_batch)
# run evaluation with the converted model
if args.run_eval:
dataset = cfg.DATASETS.TEST[0]
def __init__(self, cfg, build_model):
"""
Args:
cfg (config dict):
"""
self.data_loader = self.build_train_loader(cfg)
# Assume these objects must be constructed in this order.
model = build_model(cfg)
self.model = maybe_convert_module(model)
logger.info(f"Model: \n{self.model}")
# Assume these objects must be constructed in this order.
self.optimizer = self.build_optimizer(cfg, self.model)
if cfg.TRAINER.FP16.ENABLED:
self.mixed_precision = True
if cfg.TRAINER.FP16.TYPE == "APEX":
from apex import amp
self.model, self.optimizer = amp.initialize(
self.model,
self.optimizer,
opt_level=cfg.TRAINER.FP16.OPTS.OPT_LEVEL)
else:
self.mixed_precision = False
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
torch.cuda.set_device(comm.get_local_rank())
if cfg.MODEL.DDP_BACKEND == "torch":
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
elif cfg.MODEL.DDP_BACKEND == "apex":
from apex.parallel import DistributedDataParallel as ApexDistributedDataParallel
self.model = ApexDistributedDataParallel(self.model)
else:
raise ValueError("non-supported DDP backend: {}".format(
cfg.MODEL.DDP_BACKEND))
super().__init__(
self.model,
self.data_loader,
self.optimizer,
)
if not cfg.SOLVER.LR_SCHEDULER.get("EPOCH_WISE", False):
epoch_iters = -1
else:
epoch_iters = cfg.SOLVER.LR_SCHEDULER.get("EPOCH_ITERS")
logger.warning(f"Setup LR Scheduler in EPOCH mode: {epoch_iters}")
auto_scale_config(cfg, self.data_loader)
self.scheduler = self.build_lr_scheduler(cfg,
self.optimizer,
epoch_iters=epoch_iters)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = DefaultCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
optimizer=self.optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
self.start_epoch = 0
self.max_iter = cfg.SOLVER.LR_SCHEDULER.MAX_ITER
self.max_epoch = cfg.SOLVER.LR_SCHEDULER.MAX_EPOCH
self.window_size = cfg.TRAINER.WINDOW_SIZE
self.cfg = cfg
self.register_hooks(self.build_hooks())
class DefaultRunner(SimpleRunner):
"""
A runner with default training logic. It does the following:
1. Create a :class:`DefaultRunner` using model, optimizer, dataloader
defined by the given config. Create a LR scheduler defined by the config.
2. Load the last checkpoint or `cfg.MODEL.WEIGHTS`, if exists, when
`resume_or_load` is called.
3. Register a few common hooks defined by the config.
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:`DefaultRunner` are too much for research.
The code of this class has been annotated about restrictive assumptions it makes.
When they do not work for you, you're encouraged to:
1. Overwrite methods of this class, OR:
2. Use :class:`DefaultRunner`, 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`.
See the :doc:`/tutorials/training` tutorials for more details.
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 cvpods.
To obtain more stable behavior, write your own training logic with other public APIs.
Examples:
::
runner = DefaultRunner(cfg)
runner.resume_or_load() # load last checkpoint or MODEL.WEIGHTS
runner.train()
Attributes:
scheduler:
checkpointer (DefaultCheckpointer):
cfg (config dict):
"""
def __init__(self, cfg, build_model):
"""
Args:
cfg (config dict):
"""
self.data_loader = self.build_train_loader(cfg)
# Assume these objects must be constructed in this order.
model = build_model(cfg)
self.model = maybe_convert_module(model)
logger.info(f"Model: \n{self.model}")
# Assume these objects must be constructed in this order.
self.optimizer = self.build_optimizer(cfg, self.model)
if cfg.TRAINER.FP16.ENABLED:
self.mixed_precision = True
if cfg.TRAINER.FP16.TYPE == "APEX":
from apex import amp
self.model, self.optimizer = amp.initialize(
self.model,
self.optimizer,
opt_level=cfg.TRAINER.FP16.OPTS.OPT_LEVEL)
else:
self.mixed_precision = False
# For training, wrap with DDP. But don't need this for inference.
if comm.get_world_size() > 1:
torch.cuda.set_device(comm.get_local_rank())
if cfg.MODEL.DDP_BACKEND == "torch":
self.model = DistributedDataParallel(
self.model,
device_ids=[comm.get_local_rank()],
broadcast_buffers=False,
find_unused_parameters=True)
elif cfg.MODEL.DDP_BACKEND == "apex":
from apex.parallel import DistributedDataParallel as ApexDistributedDataParallel
self.model = ApexDistributedDataParallel(self.model)
else:
raise ValueError("non-supported DDP backend: {}".format(
cfg.MODEL.DDP_BACKEND))
super().__init__(
self.model,
self.data_loader,
self.optimizer,
)
if not cfg.SOLVER.LR_SCHEDULER.get("EPOCH_WISE", False):
epoch_iters = -1
else:
epoch_iters = cfg.SOLVER.LR_SCHEDULER.get("EPOCH_ITERS")
logger.warning(f"Setup LR Scheduler in EPOCH mode: {epoch_iters}")
auto_scale_config(cfg, self.data_loader)
self.scheduler = self.build_lr_scheduler(cfg,
self.optimizer,
epoch_iters=epoch_iters)
# Assume no other objects need to be checkpointed.
# We can later make it checkpoint the stateful hooks
self.checkpointer = DefaultCheckpointer(
# Assume you want to save checkpoints together with logs/statistics
self.model,
cfg.OUTPUT_DIR,
optimizer=self.optimizer,
scheduler=self.scheduler,
)
self.start_iter = 0
self.start_epoch = 0
self.max_iter = cfg.SOLVER.LR_SCHEDULER.MAX_ITER
self.max_epoch = cfg.SOLVER.LR_SCHEDULER.MAX_EPOCH
self.window_size = cfg.TRAINER.WINDOW_SIZE
self.cfg = cfg
self.register_hooks(self.build_hooks())
def resume_or_load(self, resume=True):
"""
If `resume==True` and `cfg.OUTPUT_DIR` contains the last checkpoint (defined by
a `last_checkpoint` file), resume from the file. Resuming means loading all
available states (eg. optimizer and scheduler) and update iteration counter
from the checkpoint. ``cfg.MODEL.WEIGHTS`` will not be used.
Otherwise, this is considered as an independent training. The method will load model
weights from the file `cfg.MODEL.WEIGHTS` (but will not load other states) and start
from iteration 0.
Args:
resume (bool): whether to do resume or not
"""
self.checkpointer.resume = resume
# 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)
if self.max_epoch is not None:
if isinstance(self.data_loader.sampler, Infinite):
length = len(self.data_loader.sampler.sampler)
else:
length = len(self.data_loader)
self.start_epoch = self.start_iter // length
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
# cfg.DATALOADER.NUM_WORKERS = 0 # save some memory and time for PreciseBN
ret = [
hooks.OptimizationHook(
accumulate_grad_steps=cfg.SOLVER.BATCH_SUBDIVISIONS,
grad_clipper=None,
mixed_precision=cfg.TRAINER.FP16.ENABLED),
hooks.LRScheduler(self.optimizer, self.scheduler),
hooks.IterationTimer(),
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,
max_iter=self.max_iter,
max_epoch=self.max_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():
# Here the default print/log frequency of each writer is used.
# run writers in the end, so that evaluation metrics are written
ret.append(
hooks.PeriodicWriter(self.build_writers(),
period=self.cfg.GLOBAL.LOG_INTERVAL))
# Put `PeriodicDumpLog` after writers so that can dump all the files,
# including the files generated by writers
return ret
def build_writers(self):
"""
Build a list of :class:`EventWriter` to be used.
It now consists of a :class:`CommonMetricPrinter`,
:class:`TensorboardXWriter` and :class:`JSONWriter`.
Args:
output_dir: directory to store JSON metrics and tensorboard events
max_iter: the total number of iterations
Returns:
list[EventWriter]: a list of :class:`EventWriter` objects.
"""
return [
# It may not always print what you want to see, since it prints "common" metrics only.
CommonMetricPrinter(
self.max_iter,
window_size=self.window_size,
epoch=self.max_epoch,
),
JSONWriter(os.path.join(self.cfg.OUTPUT_DIR, "metrics.json"),
window_size=self.window_size),
TensorboardXWriter(self.cfg.OUTPUT_DIR,
window_size=self.window_size),
]
def train(self):
"""
Run training.
Returns:
OrderedDict of results, if evaluation is enabled. Otherwise None.
"""
if self.max_epoch is None:
logger.info("Starting training from iteration {}".format(
self.start_iter))
else:
logger.info("Starting training from epoch {}".format(
self.start_epoch))
super().train(self.start_iter, self.start_epoch, self.max_iter)
if len(self.cfg.TEST.EXPECTED_RESULTS) and comm.is_main_process():
assert hasattr(self, "_last_eval_results"
), "No evaluation results obtained during training!"
verify_results(self.cfg, self._last_eval_results)
return self._last_eval_results
@classmethod
def build_optimizer(cls, cfg, model):
"""
Returns:
torch.optim.Optimizer:
It now calls :func:`cvpods.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, **kwargs):
"""
It now calls :func:`cvpods.solver.build_lr_scheduler`.
Overwrite it if you'd like a different scheduler.
"""
return build_lr_scheduler(cfg, optimizer, **kwargs)
@classmethod
def build_train_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`cvpods.data.build_train_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_train_loader(cfg)
@classmethod
def build_test_loader(cls, cfg):
"""
Returns:
iterable
It now calls :func:`cvpods.data.build_test_loader`.
Overwrite it if you'd like a different data loader.
"""
return build_test_loader(cfg)
@classmethod
def build_evaluator(cls, cfg, dataset_name):
"""
Returns:
DatasetEvaluator or None
It is not implemented by default.
"""
raise NotImplementedError(
# TODO: add this tutorial
"""
If you want DefaultRunner to automatically run evaluation,
please implement `build_evaluator()` in subclasses (see train_net.py for example).
Alternatively, you can call evaluation functions yourself (see Colab balloon tutorial for example).
""")
@classmethod
def test(cls, cfg, model, evaluators=None, output_folder=None):
"""
Args:
cfg (config dict):
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
"""
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)
# 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,
data_loader.dataset,
output_folder=output_folder)
except NotImplementedError:
logger.warn(
"No evaluator found. Use `DefaultRunner.test(evaluators=)`, "
"or implement its `build_evaluator` method.")
results[dataset_name] = {}
continue
results_i = inference_on_dataset(model, data_loader, evaluator)
if cfg.TEST.ON_FILES:
results_i = inference_on_files(evaluator)
else:
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