()
| 96 | return args |
| 97 | |
| 98 | def main(): |
| 99 | args=get_cli_args() |
| 100 | |
| 101 | mp.set_start_method(args.start_method) |
| 102 | torch.backends.cuda.matmul.allow_tf32 = True |
| 103 | try: |
| 104 | torch._dynamo.config.automatic_dynamic_shapes = False |
| 105 | except AttributeError: |
| 106 | pass |
| 107 | |
| 108 | config = K.config.load_config(args.config) |
| 109 | model_config = config['model'] |
| 110 | dataset_config = config['dataset'] |
| 111 | opt_config = config['optimizer'] |
| 112 | sched_config = config['lr_sched'] |
| 113 | ema_sched_config = config['ema_sched'] |
| 114 | cross_cond = bool(model_config['cross_cond']) |
| 115 | |
| 116 | # TODO: allow non-square input sizes |
| 117 | assert len(model_config['input_size']) == 2 and model_config['input_size'][0] == model_config['input_size'][1] |
| 118 | size = model_config['input_size'] |
| 119 | |
| 120 | accelerator = accelerate.Accelerator(gradient_accumulation_steps=args.grad_accum_steps, mixed_precision=args.mixed_precision) |
| 121 | ensure_distributed() |
| 122 | device = accelerator.device |
| 123 | unwrap = accelerator.unwrap_model |
| 124 | print(f'Process {accelerator.process_index} using device: {device}', flush=True) |
| 125 | accelerator.wait_for_everyone() |
| 126 | if accelerator.is_main_process: |
| 127 | print(f'World size: {accelerator.num_processes}', flush=True) |
| 128 | print(f'Batch size: {args.batch_size * accelerator.num_processes}', flush=True) |
| 129 | |
| 130 | if args.seed is not None: |
| 131 | seeds = torch.randint(0, 2 ** 32 - 1, [accelerator.num_processes], generator=torch.Generator().manual_seed(args.seed)) |
| 132 | print("seeds[accelerator.process_index]",seeds[accelerator.process_index]) |
| 133 | torch.manual_seed(seeds[accelerator.process_index]) |
| 134 | np.random.seed(seeds[accelerator.process_index]) |
| 135 | random.seed(seeds[accelerator.process_index]) |
| 136 | demo_gen = torch.Generator().manual_seed(torch.randint(-2 ** 63, 2 ** 63 - 1, ()).item()) |
| 137 | |
| 138 | inner_model = K.config.make_model(config) |
| 139 | inner_model_ema = deepcopy(inner_model) |
| 140 | |
| 141 | print("args.compile",args.compile) |
| 142 | if args.compile: |
| 143 | inner_model.compile() |
| 144 | # inner_model_ema.compile() |
| 145 | |
| 146 | if accelerator.is_main_process: |
| 147 | print(f'Parameters: {K.utils.n_params(inner_model):,}') |
| 148 | |
| 149 | |
| 150 | |
| 151 | use_tensorboard = args.use_tensorboard |
| 152 | if use_tensorboard and accelerator.is_main_process: |
| 153 | from torch.utils.tensorboard import SummaryWriter |
| 154 | tensorboard_writer = SummaryWriter("tensorboard_logs/"+args.name) |
| 155 |
no test coverage detected