From ad9128d1131f2ae298cee56a2393486806f23c73 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Mon, 2 Jan 2023 15:56:44 +0100 Subject: Update --- train_ti.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) (limited to 'train_ti.py') diff --git a/train_ti.py b/train_ti.py index 775b918..f1dbed1 100644 --- a/train_ti.py +++ b/train_ti.py @@ -889,9 +889,6 @@ def main(): seed=args.seed ) - if accelerator.is_main_process: - checkpointer.save_samples(global_step_offset, args.sample_steps) - local_progress_bar = tqdm( range(num_update_steps_per_epoch + num_val_steps_per_epoch), disable=not accelerator.is_local_main_process, @@ -908,6 +905,10 @@ def main(): try: for epoch in range(num_epochs): + if accelerator.is_main_process: + if epoch % args.sample_frequency == 0: + checkpointer.save_samples(global_step + global_step_offset, args.sample_steps) + local_progress_bar.set_description(f"Epoch {epoch + 1} / {num_epochs}") local_progress_bar.reset() @@ -1001,13 +1002,11 @@ def main(): "global_step": global_step + global_step_offset }) - if (epoch + 1) % args.sample_frequency == 0: - checkpointer.save_samples(global_step + global_step_offset, args.sample_steps) - # Create the pipeline using using the trained modules and save it. if accelerator.is_main_process: print("Finished! Saving final checkpoint and resume state.") checkpointer.checkpoint(global_step + global_step_offset, "end") + checkpointer.save_samples(global_step_offset, args.sample_steps) save_args(basepath, args, { "global_step": global_step + global_step_offset }) @@ -1017,6 +1016,7 @@ def main(): if accelerator.is_main_process: print("Interrupted, saving checkpoint and resume state...") checkpointer.checkpoint(global_step + global_step_offset, "end") + checkpointer.save_samples(global_step_offset, args.sample_steps) save_args(basepath, args, { "global_step": global_step + global_step_offset }) -- cgit v1.2.3-54-g00ecf