fix: do not silently scale learning rate with batch size

Also add lr_schedule parameter to configuration files
erwulff · Jun 28, 2021 · b896392 · erwulff · Jun 28, 2021 · b896392
1 parent 62a6ba3
commit b896392
Show file tree

Hide file tree

Showing 14 changed files with 15 additions and 9 deletions.
diff --git a/mlpf/tfmodel/model_setup.py b/mlpf/tfmodel/model_setup.py
@@ -554,8 +554,6 @@ def main(args, yaml_path, config):
         print("fallback to CPU", e)
         strategy = tf.distribute.OneDeviceStrategy("cpu")
         num_gpus = 0
-
-    actual_lr = global_batch_size*float(config['setup']['lr'])
 
     Xs = []
     ygens = []
@@ -580,15 +578,10 @@ def main(args, yaml_path, config):
     ygen_val = np.concatenate(ygens)
     ycand_val = np.concatenate(ycands)
 
+    lr = global_batch_size*float(config['setup']['lr'])
     with strategy.scope():
-        lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
-            actual_lr,
-            decay_steps=10000,
-            decay_rate=0.99,
-            staircase=True
-        )
         total_steps = n_epochs * n_train // global_batch_size
-        lr_schedule, optim_callbacks = get_lr_schedule(config, actual_lr, steps=total_steps)
+        lr_schedule, optim_callbacks = get_lr_schedule(config, lr, steps=total_steps)
         opt = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
         if config['setup']['dtype'] == 'float16':
             model_dtype = tf.dtypes.float16

diff --git a/parameters/cms-gnn-dense-big.yaml b/parameters/cms-gnn-dense-big.yaml
@@ -51,6 +51,7 @@ setup:
   sample_weights: inverse_sqrt
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn_dense

diff --git a/parameters/cms-gnn-dense-focal.yaml b/parameters/cms-gnn-dense-focal.yaml
@@ -52,6 +52,7 @@ setup:
   sample_weights: none
   trainable: all
   classification_loss_type: sigmoid_focal_crossentropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn_dense

diff --git a/parameters/cms-gnn-dense-transfer.yaml b/parameters/cms-gnn-dense-transfer.yaml
@@ -51,6 +51,7 @@ setup:
   sample_weights: inverse_sqrt
   trainable: transfer
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn_dense

diff --git a/parameters/cms-gnn-dense.yaml b/parameters/cms-gnn-dense.yaml
@@ -53,6 +53,7 @@ setup:
   sample_weights: inverse_sqrt
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn_dense

diff --git a/parameters/cms-gnn-skipconn-v2.yaml b/parameters/cms-gnn-skipconn-v2.yaml
@@ -51,6 +51,7 @@ setup:
   sample_weights: inverse_sqrt
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn

diff --git a/parameters/cms-gnn-skipconn.yaml b/parameters/cms-gnn-skipconn.yaml
@@ -51,6 +51,7 @@ setup:
   sample_weights: none
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn

diff --git a/parameters/cms-transformer-skipconn-gun.yaml b/parameters/cms-transformer-skipconn-gun.yaml
@@ -52,6 +52,7 @@ setup:
   sample_weights: inverse_sqrt
   trainable: all
   multi_output: yes
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: transformer

diff --git a/parameters/cms-transformer-skipconn.yaml b/parameters/cms-transformer-skipconn.yaml
@@ -50,6 +50,7 @@ setup:
   sample_weights: none
   trainable: cls
   multi_output: yes
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: transformer

diff --git a/parameters/delphes-gnn-skipconn.yaml b/parameters/delphes-gnn-skipconn.yaml
@@ -41,6 +41,7 @@ setup:
   trainable: all
   multi_output: no
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn

diff --git a/parameters/delphes-transformer-skipconn.yaml b/parameters/delphes-transformer-skipconn.yaml
@@ -39,6 +39,7 @@ setup:
   sample_weights: none
   trainable: all
   multi_output: no
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: transformer

diff --git a/parameters/test-cms-v2.yaml b/parameters/test-cms-v2.yaml
@@ -39,6 +39,7 @@ setup:
   sample_weights: none
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn_dense

diff --git a/parameters/test-cms.yaml b/parameters/test-cms.yaml
@@ -39,6 +39,7 @@ setup:
   sample_weights: none
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn

diff --git a/parameters/test-delphes.yaml b/parameters/test-delphes.yaml
@@ -38,6 +38,7 @@ setup:
   sample_weights: none
   trainable: all
   classification_loss_type: categorical_cross_entropy
+  lr_schedule: exponentialdecay  # exponentialdecay, onecycle
 
 parameters:
   model: gnn