()
| 324 | |
| 325 | |
| 326 | def save(): |
| 327 | accelerator.wait_for_everyone() |
| 328 | filename = f'{args.name}_{step:08}.pth' |
| 329 | path_checkpoints_root= os.path.join("./out_training/",args.name) |
| 330 | os.makedirs(path_checkpoints_root, exist_ok=True) |
| 331 | filename=os.path.join(path_checkpoints_root,filename) |
| 332 | if accelerator.is_main_process: |
| 333 | tqdm.write(f'Saving to {filename}...') |
| 334 | inner_model = unwrap(model.inner_model) |
| 335 | inner_model_ema = unwrap(model_ema.inner_model) |
| 336 | obj = { |
| 337 | 'config': config, |
| 338 | 'model': inner_model.state_dict(), |
| 339 | 'model_ema': inner_model_ema.state_dict(), |
| 340 | 'opt': opt.state_dict(), |
| 341 | 'sched': sched.state_dict(), |
| 342 | 'ema_sched': ema_sched.state_dict(), |
| 343 | 'epoch': epoch, |
| 344 | 'step': step, |
| 345 | 'gns_stats': gns_stats.state_dict() if gns_stats is not None else None, |
| 346 | 'ema_stats': ema_stats, |
| 347 | 'demo_gen': demo_gen.get_state(), |
| 348 | } |
| 349 | accelerator.save(obj, filename) |
| 350 | if accelerator.is_main_process: |
| 351 | state_obj = {'latest_checkpoint': filename} |
| 352 | json.dump(state_obj, open(state_path, 'w')) |
| 353 | # if args.wandb_save_model and use_wandb: |
| 354 | # wandb.save(filename) |
| 355 | #save config |
| 356 | config_path = os.path.join(path_checkpoints_root,"config.json") |
| 357 | with open(config_path, 'w', encoding='utf-8') as f: |
| 358 | json.dump(config, f, ensure_ascii=False, indent=4) |
| 359 | |
| 360 | |
| 361 | losses_since_last_print = [] |
no test coverage detected