Spaces:
Running
Running
| import math | |
| import argparse | |
| import pprint | |
| from distutils.util import strtobool | |
| from pathlib import Path | |
| from loguru import logger as loguru_logger | |
| import pytorch_lightning as pl | |
| from pytorch_lightning.utilities import rank_zero_only | |
| from pytorch_lightning.loggers import TensorBoardLogger | |
| from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor | |
| from pytorch_lightning.plugins import DDPPlugin | |
| from src.config.default import get_cfg_defaults | |
| from src.utils.misc import get_rank_zero_only_logger, setup_gpus | |
| from src.utils.profiler import build_profiler | |
| from src.lightning_trainer.data import MultiSceneDataModule | |
| from src.lightning_trainer.trainer import PL_Trainer | |
| loguru_logger = get_rank_zero_only_logger(loguru_logger) | |
| def parse_args(): | |
| # init a costum parser which will be added into pl.Trainer parser | |
| # check documentation: https://pytorch-lightning.readthedocs.io/en/latest/common/trainer.html#trainer-flags | |
| parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) | |
| parser.add_argument( | |
| 'data_cfg_path', type=str, help='data config path') | |
| parser.add_argument( | |
| 'main_cfg_path', type=str, help='main config path') | |
| parser.add_argument( | |
| '--exp_name', type=str, default='default_exp_name') | |
| parser.add_argument( | |
| '--batch_size', type=int, default=4, help='batch_size per gpu') | |
| parser.add_argument( | |
| '--num_workers', type=int, default=4) | |
| parser.add_argument( | |
| '--pin_memory', type=lambda x: bool(strtobool(x)), | |
| nargs='?', default=True, help='whether loading data to pinned memory or not') | |
| parser.add_argument( | |
| '--ckpt_path', type=str, default=None, | |
| help='pretrained checkpoint path, helpful for using a pre-trained coarse-only LoFTR') | |
| parser.add_argument( | |
| '--disable_ckpt', action='store_true', | |
| help='disable checkpoint saving (useful for debugging).') | |
| parser.add_argument( | |
| '--profiler_name', type=str, default=None, | |
| help='options: [inference, pytorch], or leave it unset') | |
| parser.add_argument( | |
| '--parallel_load_data', action='store_true', | |
| help='load datasets in with multiple processes.') | |
| parser = pl.Trainer.add_argparse_args(parser) | |
| return parser.parse_args() | |
| def main(): | |
| # parse arguments | |
| args = parse_args() | |
| rank_zero_only(pprint.pprint)(vars(args)) | |
| # init default-cfg and merge it with the main- and data-cfg | |
| config = get_cfg_defaults() | |
| config.merge_from_file(args.main_cfg_path) | |
| config.merge_from_file(args.data_cfg_path) | |
| pl.seed_everything(config.TRAINER.SEED) # reproducibility | |
| # TODO: Use different seeds for each dataloader workers | |
| # This is needed for data augmentation | |
| # scale lr and warmup-step automatically | |
| args.gpus = _n_gpus = setup_gpus(args.gpus) | |
| config.TRAINER.WORLD_SIZE = _n_gpus * args.num_nodes | |
| config.TRAINER.TRUE_BATCH_SIZE = config.TRAINER.WORLD_SIZE * args.batch_size | |
| _scaling = config.TRAINER.TRUE_BATCH_SIZE / config.TRAINER.CANONICAL_BS | |
| config.TRAINER.SCALING = _scaling | |
| config.TRAINER.TRUE_LR = config.TRAINER.CANONICAL_LR * _scaling | |
| config.TRAINER.WARMUP_STEP = math.floor(config.TRAINER.WARMUP_STEP / _scaling) | |
| # lightning module | |
| profiler = build_profiler(args.profiler_name) | |
| model = PL_Trainer(config, pretrained_ckpt=args.ckpt_path, profiler=profiler) | |
| loguru_logger.info(f"Model LightningModule initialized!") | |
| # lightning data | |
| data_module = MultiSceneDataModule(args, config) | |
| loguru_logger.info(f"Model DataModule initialized!") | |
| # TensorBoard Logger | |
| logger = TensorBoardLogger(save_dir='logs/tb_logs', name=args.exp_name, default_hp_metric=False) | |
| ckpt_dir = Path(logger.log_dir) / 'checkpoints' | |
| # Callbacks | |
| # TODO: update ModelCheckpoint to monitor multiple metrics | |
| ckpt_callback = ModelCheckpoint(monitor='auc@10', verbose=True, save_top_k=5, mode='max', | |
| save_last=True, | |
| dirpath=str(ckpt_dir), | |
| filename='{epoch}-{auc@5:.3f}-{auc@10:.3f}-{auc@20:.3f}') | |
| lr_monitor = LearningRateMonitor(logging_interval='step') | |
| callbacks = [lr_monitor] | |
| if not args.disable_ckpt: | |
| callbacks.append(ckpt_callback) | |
| # Lightning Trainer | |
| trainer = pl.Trainer.from_argparse_args( | |
| args, | |
| plugins=DDPPlugin(find_unused_parameters=False, | |
| num_nodes=args.num_nodes, | |
| sync_batchnorm=config.TRAINER.WORLD_SIZE > 0), | |
| gradient_clip_val=config.TRAINER.GRADIENT_CLIPPING, | |
| callbacks=callbacks, | |
| logger=logger, | |
| sync_batchnorm=config.TRAINER.WORLD_SIZE > 0, | |
| replace_sampler_ddp=False, # use custom sampler | |
| reload_dataloaders_every_epoch=False, # avoid repeated samples! | |
| weights_summary='full', | |
| profiler=profiler) | |
| loguru_logger.info(f"Trainer initialized!") | |
| loguru_logger.info(f"Start training!") | |
| trainer.fit(model, datamodule=data_module) | |
| if __name__ == '__main__': | |
| main() | |