Use exponentialLR

train_model.py

@@ -120,11 +120,7 @@ def train_model(batch_size=4, total_images=-1, epochs=100, scheduler_num_timeste
     text_encoder.requires_grad_(False)
 
     optimizer = torch.optim.AdamW(unet.parameters(), lr=start_learning_rate)
-    lr_scheduler = get_cosine_schedule_with_warmup(
-        optimizer=optimizer,
-        num_warmup_steps=lr_warmup_steps,
-        num_training_steps=num_images * epochs
-    )
+    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.999, verbose=True)
     model = RCTDiffusionPipeline(unet, scheduler, vae, tokenizer, text_encoder, vae_image_processor)
     unet = unet.to('cuda')
     
@@ -173,7 +169,7 @@ def train_model(batch_size=4, total_images=-1, epochs=100, scheduler_num_timeste
         if (epoch + 1) % save_model_interval == 0:
             # inference in float16
             model = RCTDiffusionPipeline(unet.to(dtype=torch.float16), scheduler, \
-                                         vae.to(dtype=torch.float16), tokenizer, text_encoder.to(dtype=torch.float16))
+                                         vae.to(dtype=torch.float16), tokenizer, text_encoder.to(dtype=torch.float16), vae_image_processor)
             save_and_test(model, epoch)
 
             # training in float32
@@ -185,4 +181,4 @@ def train_model(batch_size=4, total_images=-1, epochs=100, scheduler_num_timeste
 
 
 if __name__ == '__main__':
-    train_model(batch_size=1, total_images=4, save_model_interval=25, epochs=500, start_learning_rate=1e-5)
+    train_model(batch_size=48, save_model_interval=25, epochs=1000, start_learning_rate=1e-3)