Spaces:
Runtime error
Runtime error
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
| r""" | |
| Basic training script for PyTorch | |
| """ | |
| # Set up custom environment before nearly anything else is imported | |
| # NOTE: this should be the first import (no not reorder) | |
| from maskrcnn_benchmark.utils.env import setup_environment # noqa F401 isort:skip | |
| import argparse | |
| import os | |
| import torch | |
| from maskrcnn_benchmark.config import cfg | |
| from maskrcnn_benchmark.data import make_data_loader | |
| from maskrcnn_benchmark.solver import make_lr_scheduler | |
| from maskrcnn_benchmark.solver import make_optimizer | |
| from maskrcnn_benchmark.engine.trainer import do_train | |
| from maskrcnn_benchmark.modeling.detector import build_detection_model | |
| from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer | |
| from maskrcnn_benchmark.utils.collect_env import collect_env_info | |
| from maskrcnn_benchmark.utils.comm import synchronize, get_rank | |
| from maskrcnn_benchmark.utils.imports import import_file | |
| from maskrcnn_benchmark.utils.logging import setup_logger, Logger | |
| from maskrcnn_benchmark.utils.miscellaneous import mkdir | |
| # See if we can use apex.DistributedDataParallel instead of the torch default, | |
| # and enable mixed-precision via apex.amp | |
| try: | |
| from apex import amp | |
| except ImportError: | |
| raise ImportError('Use APEX for multi-precision via apex.amp') | |
| def train(cfg, local_rank, distributed): | |
| model = build_detection_model(cfg) | |
| device = torch.device(cfg.MODEL.DEVICE) | |
| model.to(device) | |
| optimizer = make_optimizer(cfg, model) | |
| scheduler = make_lr_scheduler(cfg, optimizer) | |
| # Initialize mixed-precision training | |
| use_mixed_precision = cfg.DTYPE == "float16" | |
| amp_opt_level = 'O1' if use_mixed_precision else 'O0' | |
| model, optimizer = amp.initialize(model, optimizer, opt_level=amp_opt_level) | |
| if distributed: | |
| model = torch.nn.parallel.DistributedDataParallel( | |
| model, device_ids=[local_rank], output_device=local_rank, | |
| # this should be removed if we update BatchNorm stats | |
| broadcast_buffers=False, | |
| # find_unused_parameters=True | |
| ) | |
| arguments = {} | |
| arguments["iteration"] = 0 | |
| output_dir = cfg.OUTPUT_DIR | |
| save_to_disk = get_rank() == 0 | |
| checkpointer = DetectronCheckpointer( | |
| cfg, model, optimizer, scheduler, output_dir, save_to_disk | |
| ) | |
| extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT, resume=cfg.SOLVER.RESUME) | |
| if cfg.SOLVER.RESUME: | |
| arguments.update(extra_checkpoint_data) | |
| data_loader = make_data_loader( | |
| cfg, | |
| is_train=True, | |
| is_distributed=distributed, | |
| start_iter=arguments["iteration"], | |
| ) | |
| checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD | |
| tb_logger = Logger(cfg.OUTPUT_DIR, local_rank) | |
| do_train( | |
| model, | |
| data_loader, | |
| optimizer, | |
| scheduler, | |
| checkpointer, | |
| device, | |
| checkpoint_period, | |
| arguments, | |
| tb_logger, | |
| cfg, | |
| local_rank, | |
| ) | |
| return model | |
| def main(): | |
| parser = argparse.ArgumentParser(description="PyTorch Object Detection Training") | |
| parser.add_argument( | |
| "--config-file", | |
| default="", | |
| metavar="FILE", | |
| help="path to config file", | |
| type=str, | |
| ) | |
| parser.add_argument("--local_rank", type=int, default=0) | |
| parser.add_argument( | |
| "--skip-test", | |
| dest="skip_test", | |
| help="Do not test the final model", | |
| action="store_true", | |
| ) | |
| parser.add_argument( | |
| "opts", | |
| help="Modify config options using the command-line", | |
| default=None, | |
| nargs=argparse.REMAINDER, | |
| ) | |
| args = parser.parse_args() | |
| num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 | |
| args.distributed = num_gpus > 1 | |
| if args.distributed: | |
| torch.cuda.set_device(args.local_rank) | |
| torch.distributed.init_process_group( | |
| backend="nccl", init_method="env://" | |
| ) | |
| cfg.merge_from_file(args.config_file) | |
| cfg.merge_from_list(args.opts) | |
| cfg.freeze() | |
| output_dir = cfg.OUTPUT_DIR | |
| if output_dir: | |
| mkdir(output_dir) | |
| local_rank = get_rank() | |
| logger = setup_logger("maskrcnn_benchmark", output_dir, local_rank) | |
| if local_rank == 0: | |
| logger.info("Using {} GPUs".format(num_gpus)) | |
| logger.info(args) | |
| logger.info("Collecting env info (might take some time)") | |
| logger.info("\n" + collect_env_info()) | |
| logger.info("Loaded configuration file {}".format(args.config_file)) | |
| with open(args.config_file, "r") as cf: | |
| config_str = "\n" + cf.read() | |
| logger.info(config_str) | |
| logger.info("Running with config:\n{}".format(cfg)) | |
| model = train(cfg, args.local_rank, args.distributed) | |
| if __name__ == "__main__": | |
| main() | |