fix schedulefree saving, warmup_steps

feffy380 · Apr 11, 2024 · 198096b · 198096b
1 parent 8f6d4d7
commit 198096b
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/train_network.py b/train_network.py
@@ -450,6 +450,18 @@ def train(self, args):
         # 学習に必要なクラスを準備する
         accelerator.print("prepare optimizer, data loader etc.")
 
+        schedulefree = "schedulefree" in args.optimizer_type.lower()
+        if schedulefree:
+            args.lr_scheduler = "constant"
+            if args.lr_warmup_steps:
+                if args.optimizer_args is None:
+                    args.optimizer_args = []
+                warmup_steps = args.lr_warmup_steps
+                if warmup_steps < 1.0:
+                    warmup_steps = int(warmup_steps * args.max_train_steps)
+                args.optimizer_args.append(f"warmup_steps={warmup_steps}")
+                args.lr_warmup_steps = None
+
         # 後方互換性を確保するよ
         try:
             trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, args.learning_rate)
@@ -468,7 +480,6 @@ def train(self, args):
             trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)
 
         optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
-        schedulefree = "schedulefree" in args.optimizer_type.lower()
 
         # dataloaderを準備する
         # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
@@ -497,9 +508,6 @@ def train(self, args):
         train_dataset_group.set_max_train_steps(args.max_train_steps)
 
         # lr schedulerを用意する
-        if "schedulefree" in args.optimizer_type.lower():
-            args.lr_scheduler = "constant"
-            args.lr_warmup_steps = None
         lr_scheduler = train_util.get_scheduler_fix(args, optimizer, accelerator.num_processes)
 
         # 実験的機能：勾配も含めたfp16/bf16学習を行う　モデル全体をfp16/bf16にする
@@ -1015,14 +1023,14 @@ def update_embeddings_map(accelerator, text_encoders, embeddings_map, embedding_
             if args.continue_inversion:
                 for t_enc in text_encoders:
                     t_enc.train()
-            if schedulefree:
-                optimizer.optimizer.train()
 
             metadata["ss_epoch"] = str(epoch + 1)
 
             accelerator.unwrap_model(network).on_epoch_start(text_encoder, unet)
 
             for step, batch in enumerate(train_dataloader):
+                if schedulefree:
+                    optimizer.optimizer.train()
                 current_step.value = global_step
                 with accelerator.accumulate(training_model):
                     on_step_start(text_encoder, unet)

diff --git a/train_textual_inversion.py b/train_textual_inversion.py
@@ -406,11 +406,23 @@ def train(self, args):
 
         # 学習に必要なクラスを準備する
         accelerator.print("prepare optimizer, data loader etc.")
+
+        schedulefree = "schedulefree" in args.optimizer_type.lower()
+        if schedulefree:
+            args.lr_scheduler = "constant"
+            if args.lr_warmup_steps:
+                if args.optimizer_args is None:
+                    args.optimizer_args = []
+                warmup_steps = args.lr_warmup_steps
+                if warmup_steps < 1.0:
+                    warmup_steps = int(warmup_steps * args.max_train_steps)
+                args.optimizer_args.append(f"warmup_steps={warmup_steps}")
+                args.lr_warmup_steps = None
+
         trainable_params = []
         for text_encoder in text_encoders:
             trainable_params += text_encoder.get_input_embeddings().parameters()
         _, _, optimizer = train_util.get_optimizer(args, trainable_params)
-        schedulefree = "schedulefree" in args.optimizer_type.lower()
 
         # dataloaderを準備する
         # DataLoaderのプロセス数：0 は persistent_workers が使えないので注意
@@ -583,12 +595,12 @@ def remove_model(old_ckpt_name):
 
             for text_encoder in text_encoders:
                 text_encoder.train()
-            if schedulefree:
-                optimizer.optimizer.train()
 
             loss_total = 0
 
             for step, batch in enumerate(train_dataloader):
+                if schedulefree:
+                    optimizer.optimizer.train()
                 current_step.value = global_step
                 with accelerator.accumulate(text_encoders[0]):
                     # disable ToDo after some steps