1 files changed, 8 insertions, 10 deletions
diff --git a/config.py b/config.py
index 4c108e2..e70c2bd 100644
--- a/config.py
+++ b/config.py
@@ -72,8 +72,6 @@ config: Configuration = {
         },
         'optimizer': {
             # Global parameters
-            # Iteration start to optimize non-disentangling parts
-            # 'start_iter': 0,
             # Initial learning rate of Adam Optimizer
             'lr': 1e-4,
             # Coefficients used for computing running averages of
@@ -87,15 +85,15 @@ config: Configuration = {
             # 'amsgrad': False,
 
             # Local parameters (override global ones)
-            # 'auto_encoder': {
-            #     'weight_decay': 0.001
-            # },
+            'auto_encoder': {
+                'weight_decay': 0.001
+            },
         },
         'scheduler': {
-            # Period of learning rate decay
-            'step_size': 500,
-            # Multiplicative factor of decay
-            'gamma': 1,
+            # Step start to decay
+            'start_step': 15_000,
+            # Multiplicative factor of decay in the end
+            'final_gamma': 0.001,
         }
     },
     # Model metadata
@@ -109,6 +107,6 @@ config: Configuration = {
         # Restoration iteration (multiple models, e.g. nm, bg and cl)
         'restore_iters': (0, 0, 0),
         # Total iteration for training (multiple models)
-        'total_iters': (80_000, 80_000, 80_000),
+        'total_iters': (25_000, 25_000, 25_000),
     },
 }