| 46 | dist.init_process_group(world_size=1, rank=0, store=dist.HashStore()) |
| 47 | |
| 48 | def get_cli_args(): |
| 49 | p = argparse.ArgumentParser(description=__doc__, |
| 50 | formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| 51 | p.add_argument('--dataset_path', required=True, help='Path to the difflocks dataset to train on') |
| 52 | p.add_argument('--dataset_processed_path', required=True, help='Path to the difflocks processed dataset to train on') |
| 53 | p.add_argument('--batch-size', type=int, default=8, |
| 54 | help='the batch size') |
| 55 | p.add_argument('--checkpointing', action='store_true', |
| 56 | help='enable gradient checkpointing') |
| 57 | p.add_argument('--compile', action='store_true', |
| 58 | help='compile the model') |
| 59 | p.add_argument('--config', type=str, required=True, |
| 60 | help='the configuration file') |
| 61 | p.add_argument('--demo-every', type=int, default=500, |
| 62 | help='save a demo grid every this many steps') |
| 63 | p.add_argument('--end-step', type=int, default=None, |
| 64 | help='the step to end training at') |
| 65 | p.add_argument('--gns', action='store_true', |
| 66 | help='measure the gradient noise scale (DDP only, disables stratified sampling)') |
| 67 | p.add_argument('--grad-accum-steps', type=int, default=1, |
| 68 | help='the number of gradient accumulation steps') |
| 69 | p.add_argument('--lr', type=float, |
| 70 | help='the learning rate') |
| 71 | p.add_argument('--mixed-precision', type=str, |
| 72 | help='the mixed precision type') |
| 73 | p.add_argument('--name', type=str, default='model', |
| 74 | help='the name of the run') |
| 75 | p.add_argument('--num-workers', type=int, default=8, |
| 76 | help='the number of data loader workers') |
| 77 | p.add_argument('--reset-ema', action='store_true', |
| 78 | help='reset the EMA') |
| 79 | p.add_argument('--resume', type=str, |
| 80 | help='the checkpoint to resume from') |
| 81 | p.add_argument('--resume-inference', type=str, |
| 82 | help='the inference checkpoint to resume from') |
| 83 | p.add_argument('--save-checkpoints', action='store_true', |
| 84 | help='save checkpoints every save-every stps') |
| 85 | p.add_argument('--save-every', type=int, default=10000, |
| 86 | help='save every this many steps') |
| 87 | p.add_argument('--seed', type=int, default=0, |
| 88 | help='the random seed') |
| 89 | p.add_argument('--start-method', type=str, default='spawn', |
| 90 | choices=['fork', 'forkserver', 'spawn'], |
| 91 | help='the multiprocessing start method') |
| 92 | p.add_argument('--use-tensorboard', action='store_true', |
| 93 | help='flag to tuse tensorboard for logging scalars and images') |
| 94 | args = p.parse_args() |
| 95 | |
| 96 | return args |
| 97 | |
| 98 | def main(): |
| 99 | args=get_cli_args() |