Add constant lr scheduler

crowsonkb · Feb 6, 2023 · 6e90ce8 · 6e90ce8
1 parent 6a151de
commit 6e90ce8
Show file tree

Hide file tree

Showing 6 changed files with 7 additions and 20 deletions.
diff --git a/configs/config_32x32_small.json b/configs/config_32x32_small.json
@@ -33,10 +33,7 @@
         "weight_decay": 1e-3
     },
     "lr_sched": {
-        "type": "inverse",
-        "inv_gamma": 20000.0,
-        "power": 1.0,
-        "warmup": 0.99
+        "type": "constant"
     },
     "ema_sched": {
         "type": "inverse",

diff --git a/configs/config_32x32_small_butterflies.json b/configs/config_32x32_small_butterflies.json
@@ -34,10 +34,7 @@
         "weight_decay": 1e-3
     },
     "lr_sched": {
-        "type": "inverse",
-        "inv_gamma": 20000.0,
-        "power": 1.0,
-        "warmup": 0.99
+        "type": "constant"
     },
     "ema_sched": {
         "type": "inverse",

diff --git a/configs/config_cifar10.json b/configs/config_cifar10.json
@@ -33,10 +33,7 @@
         "weight_decay": 1e-3
     },
     "lr_sched": {
-        "type": "inverse",
-        "inv_gamma": 20000.0,
-        "power": 1.0,
-        "warmup": 0.99
+        "type": "constant"
     },
     "ema_sched": {
         "type": "inverse",

diff --git a/configs/config_mnist.json b/configs/config_mnist.json
@@ -33,10 +33,7 @@
         "weight_decay": 1e-3
     },
     "lr_sched": {
-        "type": "inverse",
-        "inv_gamma": 20000.0,
-        "power": 1.0,
-        "warmup": 0.99
+        "type": "constant"
     },
     "ema_sched": {
         "type": "inverse",

diff --git a/k_diffusion/config.py b/k_diffusion/config.py
@@ -34,10 +34,7 @@ def load_config(file):
             'weight_decay': 1e-3,
         },
         'lr_sched': {
-            'type': 'inverse',
-            'inv_gamma': 20000.,
-            'power': 1.,
-            'warmup': 0.99,
+            'type': 'constant',
         },
         'ema_sched': {
             'type': 'inverse',

diff --git a/train.py b/train.py
@@ -130,6 +130,8 @@ def main():
                                       num_steps=sched_config['num_steps'],
                                       decay=sched_config['decay'],
                                       warmup=sched_config['warmup'])
+    elif sched_config['type'] == 'constant':
+        sched = optim.lr_scheduler.LambdaLR(opt, lambda _: 1.0)
     else:
         raise ValueError('Invalid schedule type')