Spaces:
Sleeping
Sleeping
| """ | |
| DEIM: DETR with Improved Matching for Fast Convergence | |
| Copyright (c) 2024 The DEIM Authors. All Rights Reserved. | |
| --------------------------------------------------------------------------------- | |
| Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR) | |
| Copyright (c) 2023 lyuwenyu. All Rights Reserved. | |
| """ | |
| import os | |
| import sys | |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..')) | |
| import argparse | |
| from engine.misc import dist_utils | |
| from engine.core import YAMLConfig, yaml_utils | |
| from engine.solver import TASKS | |
| try: | |
| import wandb | |
| _WANDB_AVAILABLE = True | |
| except ImportError: | |
| _WANDB_AVAILABLE = False | |
| wandb = None | |
| debug=False | |
| if debug: | |
| import torch | |
| def custom_repr(self): | |
| return f'{{Tensor:{tuple(self.shape)}}} {original_repr(self)}' | |
| original_repr = torch.Tensor.__repr__ | |
| torch.Tensor.__repr__ = custom_repr | |
| def main(args, ) -> None: | |
| """main | |
| """ | |
| dist_utils.setup_distributed(args.print_rank, args.print_method, seed=args.seed) | |
| assert not all([args.tuning, args.resume]), \ | |
| 'Only support from_scrach or resume or tuning at one time' | |
| update_dict = yaml_utils.parse_cli(args.update) | |
| update_dict.update({k: v for k, v in args.__dict__.items() \ | |
| if k not in ['update', ] and v is not None}) | |
| cfg = YAMLConfig(args.config, **update_dict) | |
| if args.resume or args.tuning: | |
| if 'HGNetv2' in cfg.yaml_cfg: | |
| cfg.yaml_cfg['HGNetv2']['pretrained'] = False | |
| # Initialize wandb if enabled and on main process | |
| wandb_run = None | |
| if _WANDB_AVAILABLE and args.use_wandb and dist_utils.is_main_process(): | |
| # Override wandb enabled setting from CLI | |
| if not cfg.wandb: | |
| cfg.wandb = {} | |
| cfg.wandb['enabled'] = True | |
| if args.wandb_id: | |
| cfg.wandb['id'] = args.wandb_id | |
| wandb_run = wandb.init( | |
| project=cfg.wandb.get('project', 'DEIM-detection'), | |
| entity=cfg.wandb.get('entity', None), | |
| name=cfg.wandb.get('name', None), | |
| id=cfg.wandb.get('id', None), | |
| config=cfg.yaml_cfg, | |
| resume='auto' if cfg.wandb.get('id') else False, | |
| reinit=True, | |
| tags=cfg.wandb.get('tags', []), | |
| notes=cfg.wandb.get('notes', None), | |
| group=cfg.wandb.get('group', None), | |
| job_type=cfg.wandb.get('job_type', 'training'), | |
| anonymous=cfg.wandb.get('anonymous', None), | |
| mode=cfg.wandb.get('mode', 'online'), | |
| save_code=cfg.wandb.get('save_code', True), | |
| ) | |
| cfg.wandb_run = wandb_run | |
| print("WandB initialized successfully!") | |
| elif args.use_wandb and not _WANDB_AVAILABLE: | |
| print("Warning: --use-wandb specified but wandb is not available. Please install wandb.") | |
| elif args.use_wandb and not dist_utils.is_main_process(): | |
| print("WandB will only log from main process in distributed training.") | |
| print('cfg: ', cfg.__dict__) | |
| solver = TASKS[cfg.yaml_cfg['task']](cfg) | |
| if args.test_only: | |
| solver.val() | |
| else: | |
| solver.fit() | |
| # Finish wandb run | |
| if wandb_run is not None: | |
| wandb_run.finish() | |
| dist_utils.cleanup() | |
| if __name__ == '__main__': | |
| parser = argparse.ArgumentParser() | |
| # priority 0 | |
| parser.add_argument('-c', '--config', type=str, required=True) | |
| parser.add_argument('-r', '--resume', type=str, help='resume from checkpoint') | |
| parser.add_argument('-t', '--tuning', type=str, help='tuning from checkpoint') | |
| parser.add_argument('-d', '--device', type=str, help='device',) | |
| parser.add_argument('--seed', type=int, help='exp reproducibility') | |
| parser.add_argument('--use-amp', action='store_true', help='auto mixed precision training') | |
| parser.add_argument('--output-dir', type=str, help='output directoy') | |
| parser.add_argument('--summary-dir', type=str, help='tensorboard summry') | |
| parser.add_argument('--test-only', action='store_true', default=False,) | |
| # wandb arguments | |
| parser.add_argument('--use-wandb', action='store_true', help='enable wandb logging') | |
| parser.add_argument('--wandb-id', type=str, help='resume existing run id') | |
| parser.add_argument('--enable-vis', action='store_true', help='enable visualization of augmented images during training') | |
| # priority 1 | |
| parser.add_argument('-u', '--update', nargs='+', help='update yaml config') | |
| # env | |
| parser.add_argument('--print-method', type=str, default='builtin', help='print method') | |
| parser.add_argument('--print-rank', type=int, default=0, help='print rank id') | |
| parser.add_argument('--local-rank', type=int, help='local rank id') | |
| args = parser.parse_args() | |
| main(args) | |