Merge branch 'master' of github.com:experimaestro/experimaestro-ir

experimaestro · Jul 18, 2023 · 7750213 · 7750213
2 parents 0adc839 + a5575e0
commit 7750213
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 3 deletions.
diff --git a/src/xpmir/papers/helpers/optim.py b/src/xpmir/papers/helpers/optim.py
@@ -54,3 +54,25 @@ def optimizer(self):
                 ),
             ]
         )
+
+    @cached_property
+    def optimizer_splade(self):
+        scheduler = (
+            LinearWithWarmup(
+                num_warmup_steps=self.num_warmup_steps,
+                min_factor=self.warmup_min_factor,
+            )
+            if self.scheduler
+            else None
+        )
+
+        return get_optimizers(
+            [
+                ParameterOptimizer(
+                    scheduler=scheduler,
+                    optimizer=AdamW(
+                        lr=self.lr, weight_decay=self.weight_decay, eps=self.eps
+                    ),
+                ),
+            ]
+        )
diff --git a/src/xpmir/papers/splade/experiment.py b/src/xpmir/papers/splade/experiment.py
@@ -177,7 +177,7 @@ def run(
         # the model to be trained
         model=spladev2,
         # Optimization settings
-        optimizers=cfg.splade.optimization.optimizer,
+        optimizers=cfg.splade.optimization.optimizer_splade,
         steps_per_epoch=cfg.splade.optimization.steps_per_epoch,
         use_fp16=True,
         max_epochs=cfg.splade.optimization.max_epochs,

diff --git a/src/xpmir/papers/splade/normal.yaml b/src/xpmir/papers/splade/normal.yaml
@@ -24,11 +24,11 @@ splade:
 
     model: splade_max
     # validation for each (steps_per_epoch * validation interval) steps
-    validation_interval: 40
+    validation_interval: 8
     lambda_q: 3.0e-4
     lambda_d: 1.0e-4
     lamdba_warmup_steps: 50000
-    requirements: duration=2 days & cuda(mem=40G) * 2
+    requirements: duration=4 days & cuda(mem=40G)
 
 retrieval:
     requirements: duration=2 days & cuda(mem=24G)