update save format (#28)

* fix

* update settings

* Update dtype to torch.float32 in DiT model

* Remove unnecessary assert statement in DiT class

* Refactor logging and checkpoint saving in train_img.py

* Update save directory to include global step in ckpt_utils.py and log global step in train_img.py

opendit/utils/ckpt_utils.py

@@ -57,12 +57,13 @@ def save(
     lr_scheduler: _LRScheduler,
     epoch: int,
     step: int,
+    global_step: int,
     batch_size: int,
     coordinator: DistCoordinator,
     save_dir: str,
     shape_dict: dict,
 ):
-    save_dir = os.path.join(save_dir, f"epoch{epoch}-step{step}")
+    save_dir = os.path.join(save_dir, f"epoch{epoch}-global_step{global_step}")
     os.makedirs(os.path.join(save_dir, "model"), exist_ok=True)
 
     booster.save_model(model, os.path.join(save_dir, "model"), shard=True)
@@ -79,6 +80,7 @@ def save(
     running_states = {
         "epoch": epoch,
         "step": step,
+        "global_step": global_step,
         "sample_start_index": step * batch_size,
     }
     if coordinator.is_master():

train_img.py

@@ -244,12 +244,15 @@ def main(args):
                         lr_scheduler,
                         epoch,
                         step + 1,
+                        global_step + 1,
                         args.batch_size,
                         coordinator,
                         experiment_dir,
                         ema_shape_dict,
                     )
-                    logger.info(f"Saved checkpoint at epoch {epoch} step {step + 1} to {experiment_dir}")
+                    logger.info(
+                        f"Saved checkpoint at epoch {epoch} step {step + 1} global_step {global_step + 1} to {experiment_dir}"
+                    )
 
         # the continue epochs are not resumed, so we need to reset the sampler start index and start step
         dataloader.sampler.set_start_index(0)