feat: multistep lr scheduler

models/modules/utils.py

@@ -138,6 +138,10 @@ def lambda_rule(epoch):
         scheduler = lr_scheduler.StepLR(
             optimizer, step_size=opt.train_lr_decay_iters, gamma=0.1
         )
+    elif opt.train_lr_policy == "multistep":
+        scheduler = lr_scheduler.MultiStepLR(
+            optimizer, milestones=opt.train_lr_steps, gamma=0.1
+        )
     elif opt.train_lr_policy == "plateau":
         scheduler = lr_scheduler.ReduceLROnPlateau(
             optimizer, mode="min", factor=0.2, threshold=0.01, patience=5

options/train_options.py

@@ -298,7 +298,7 @@ def initialize(self, parser):
             "--train_lr_policy",
             type=str,
             default="linear",
-            choices=["linear", "step", "plateau", "cosine"],
+            choices=["linear", "step", "multistep", "plateau", "cosine"],
             help="learning rate policy.",
         )
         parser.add_argument(
@@ -307,6 +307,14 @@ def initialize(self, parser):
             default=50,
             help="multiply by a gamma every lr_decay_iters iterations",
         )
+        parser.add_argument(
+            "--train_lr_steps",
+            default=[],
+            nargs="*",
+            type=int,
+            help="number of epochs between reductions of the learning rate by gamma=0.1",
+        )
+
         parser.add_argument(
             "--train_nb_img_max_fid",
             type=int,
@@ -700,7 +708,7 @@ def _after_parse(self, opt, set_device=True):
             )
 
         # vitclip16 projector only works with input size 224
-        if opt.D_proj_network_type == "efficientnet":
+        if "projected_d" in opt.D_netDs and opt.D_proj_network_type == "efficientnet":
             if opt.D_proj_interp < 224:
                 warnings.warn(
                     "Efficiennet projector has minimal input size of 224, setting D_proj_interp to 224"