#!/usr/bin/env python # Copyright (c) Facebook, Inc. and its affiliates. """ Training script using the new "LazyConfig" python config files. This scripts reads a given python config file and runs the training or evaluation. It can be used to train any models or dataset as long as they can be instantiated by the recursive construction defined in the given config file. Besides lazy construction of models, dataloader, etc., this scripts expects a few common configuration parameters currently defined in "configs/common/train.py". To add more complicated training logic, you can easily add other configs in the config file and implement a new train_net.py to handle them. """ import argparse import logging import os import torch import torch.distributed as dist from detectron2.config import LazyConfig, instantiate from detectron2.engine import (AMPTrainer, SimpleTrainer, default_argument_parser, default_setup, default_writers, hooks, launch) from detectron2.engine.defaults import create_ddp_model from detectron2.evaluation import inference_on_dataset, print_csv_format from detectron2.utils import comm # from detectron2.checkpoint import DetectionCheckpointer from detectron2_pe.checkpoint import DetectionCheckpointer logger = logging.getLogger("detectron2") def do_test(cfg, model): if "evaluator" in cfg.dataloader: ret = inference_on_dataset( model, instantiate(cfg.dataloader.test), instantiate(cfg.dataloader.evaluator), ) print_csv_format(ret) return ret def do_train(args, cfg): """ Args: cfg: an object with the following attributes: model: instantiate to a module dataloader.{train,test}: instantiate to dataloaders dataloader.evaluator: instantiate to evaluator for test set optimizer: instantaite to an optimizer lr_multiplier: instantiate to a fvcore scheduler train: other misc config defined in `configs/common/train.py`, including: output_dir (str) init_checkpoint (str) amp.enabled (bool) max_iter (int) eval_period, log_period (int) device (str) checkpointer (dict) ddp (dict) """ model = instantiate(cfg.model) logger = logging.getLogger("detectron2") logger.info("Model:\n{}".format(model)) model.to(cfg.train.device) cfg.optimizer.params.model = model optim = instantiate(cfg.optimizer) train_loader = instantiate(cfg.dataloader.train) model = create_ddp_model(model, **cfg.train.ddp) trainer = (AMPTrainer if cfg.train.amp.enabled else SimpleTrainer)( model, train_loader, optim ) checkpointer = DetectionCheckpointer( model, cfg.train.output_dir, trainer=trainer, ) trainer.register_hooks( [ hooks.IterationTimer(), hooks.LRScheduler(scheduler=instantiate(cfg.lr_multiplier)), ( hooks.PeriodicCheckpointer(checkpointer, **cfg.train.checkpointer) if comm.is_main_process() else None ), hooks.EvalHook(cfg.train.eval_period, lambda: do_test(cfg, model)), ( hooks.PeriodicWriter( default_writers(cfg.train.output_dir, cfg.train.max_iter), period=cfg.train.log_period, ) if comm.is_main_process() else None ), ] ) checkpointer.resume_or_load(cfg.train.init_checkpoint, resume=args.resume) if args.resume and checkpointer.has_checkpoint(): # The checkpoint stores the training iteration that just finished, thus we start # at the next iteration start_iter = trainer.iter + 1 else: start_iter = 0 trainer.train(start_iter, cfg.train.max_iter) def main(args): cfg = LazyConfig.load(args.config_file) cfg = LazyConfig.apply_overrides(cfg, args.opts) default_setup(cfg, args) if args.eval_only: model = instantiate(cfg.model) model.to(cfg.train.device) model = create_ddp_model(model) DetectionCheckpointer(model).load(cfg.train.init_checkpoint) print(do_test(cfg, model)) else: do_train(args, cfg) def invoke_main() -> None: args = default_argument_parser().parse_args() launch( main, args.num_gpus, num_machines=args.num_machines, machine_rank=args.machine_rank, dist_url=args.dist_url, args=(args,), ) def default_argument_parser_slurm(epilog=None): """ Create a parser with some common arguments used by detectron2 users. Args: epilog (str): epilog passed to ArgumentParser describing the usage. Returns: argparse.ArgumentParser: """ parser = argparse.ArgumentParser( epilog=epilog or f"Launch using torchrun", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "--config-file", default="", metavar="FILE", help="path to config file" ) parser.add_argument( "--resume", action="store_true", help="Whether to attempt to resume from the checkpoint directory. " "See documentation of `DefaultTrainer.resume_or_load()` for what it means.", ) parser.add_argument( "--eval-only", action="store_true", help="perform evaluation only" ) # parser.add_argument("--num-gpus", type=int, default=1, help="number of gpus *per machine*") # parser.add_argument("--num-machines", type=int, default=1, help="total number of machines") # parser.add_argument( # "--machine-rank", type=int, default=0, help="the rank of this machine (unique per machine)" # ) # # PyTorch still may leave orphan processes in multi-gpu training. # # Therefore we use a deterministic way to obtain port, # # so that users are aware of orphan processes by seeing the port occupied. # port = 2**15 + 2**14 + hash(os.getuid() if sys.platform != "win32" else 1) % 2**14 # parser.add_argument( # "--dist-url", # default="tcp://127.0.0.1:{}".format(port), # help="initialization URL for pytorch distributed backend. See " # "https://pytorch.org/docs/stable/distributed.html for details.", # ) parser.add_argument( "opts", help=""" Modify config options at the end of the command. For Yacs configs, use space-separated "PATH.KEY VALUE" pairs. For python-based LazyConfig, use "path.key=value". """.strip(), default=None, nargs=argparse.REMAINDER, ) return parser def main_worker(main_func, args): world_size = int(os.environ["WORLD_SIZE"]) rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) has_gpu = torch.cuda.is_available() if has_gpu: torch.cuda.set_device(local_rank) dist.init_process_group( backend="nccl" if has_gpu else "gloo", init_method="env://", ) # Setup the local process group num_gpus_per_machine = torch.cuda.device_count() machine_rank = rank // num_gpus_per_machine # comm.create_local_process_group(num_gpus_per_machine) # Setup the local process group (which contains ranks within the same machine) assert comm._LOCAL_PROCESS_GROUP is None num_machines = world_size // num_gpus_per_machine for i in range(num_machines): ranks_on_i = list( range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine) ) pg = dist.new_group(ranks_on_i) if i == machine_rank: comm._LOCAL_PROCESS_GROUP = pg # synchronize to prevent possible timeout after calling init_process_group comm.synchronize() main_func(args) def slurm_main() -> None: """ Launch the main function using torchrun. This function should be called in the main script. """ args = default_argument_parser_slurm().parse_args() main_worker(main, args) if __name__ == "__main__": # invoke_main() # pragma: no cover slurm_main()