| 26 | |
| 27 | |
| 28 | def parse_args(): |
| 29 | parser = argparse.ArgumentParser(description='Training') |
| 30 | parser.add_argument('--config', |
| 31 | help='Path to the training config file.', required=True) |
| 32 | parser.add_argument('--logdir', help='Dir for saving logs and models.') |
| 33 | parser.add_argument('--checkpoint', default='', help='Checkpoint path.') |
| 34 | parser.add_argument('--seed', type=int, default=2, help='Random seed.') |
| 35 | parser.add_argument('--randomized_seed', action='store_true', help='Use a random seed between 0-10000.') |
| 36 | parser.add_argument('--local_rank', type=int, default=os.getenv('LOCAL_RANK', 0)) |
| 37 | parser.add_argument('--single_gpu', action='store_true') |
| 38 | parser.add_argument('--debug', action='store_true') |
| 39 | parser.add_argument('--use_jit', action='store_true') |
| 40 | parser.add_argument('--profile', action='store_true') |
| 41 | parser.add_argument('--wandb', action='store_true') |
| 42 | parser.add_argument('--wandb_name', default='default', type=str) |
| 43 | parser.add_argument('--wandb_id', type=str) |
| 44 | parser.add_argument('--resume', type=int) |
| 45 | parser.add_argument('--num_workers', type=int) |
| 46 | args = parser.parse_args() |
| 47 | return args |
| 48 | |
| 49 | |
| 50 | def main(): |