Collection of minor changes

config_paper/default_tagging.yaml

@@ -20,9 +20,13 @@ data:
 train: true
 training:
  iterations: 2e5
+ optimizer: Lion
  batchsize: 128
  lr: 3e-4
  validate_every_n_steps: 5000
+ weight_decay: 0 # top tagging from scratch: 0.2; JetClass pretraining/training: 0; top finetuning: 0.01
+ scheduler: CosineAnnealingLR
+
 
  force_xformers: true
 

experiments/base_experiment.py

@@ -499,8 +499,6 @@ def cycle(iterable):
         for step in range(self.cfg.training.iterations):
             # training
             self.model.train()
-            if self.cfg.training.optimizer == "ScheduleFree":
-                self.optimizer.train()
             data = next(iterator)
             self._step(data, step)
 
@@ -632,8 +630,6 @@ def _validate(self, step):
         metrics = self._init_metrics()
 
         self.model.eval()
-        if self.cfg.training.optimizer == "ScheduleFree":
-            self.optimizer.eval()
         with torch.no_grad():
             for data in self.val_loader:
                 # use EMA for validation if available

requirements.txt

@@ -5,7 +5,7 @@ black==22.3.0
 ## Standard science stack
 h5py
 matplotlib
-numpy<1.25
+numpy==1.25
 scipy
 
 ## Standard utils