Add sweep conf

1 files changed, 29 insertions, 0 deletions

diff --git a/training/conf/sweep/conv_transformer.yaml b/training/conf/sweep/conv_transformer.yaml
new file mode 100644
index 0000000..70e0a56
--- /dev/null
+++ b/training/conf/sweep/conv_transformer.yaml
@@ -0,0 +1,29 @@
+# @package _global_
+
+# Example:
+# python main.py -m sweep=lines experiment=conv_transformer_lines
+
+defaults:
+  - override /hydra/sweeper: optuna
+
+optimized_metric: train/loss
+
+hydra:
+  mode: "MULTIRUN"
+  sweeper:
+    _target_: hydra_plugins.hydra_optuna_sweeper.optuna_sweeper.OptunaSweeper
+    storage: null
+    study_name: null
+    n_jobs: 1
+    direction: minimize
+    n_trials: 20
+    sampler:
+      _target_: optuna.samplers.TPESampler
+      seed: 4711
+      n_startup_trials: 5
+
+    params:
+      # optimizer: interval(0.0003, 0.001)
+      network.decoder.decoder.depth: choice(4, 8, 10)
+      network.decoder.decoder.block.self_attn.num_heads: choice(4, 6, 8)
+      network.decoder.decoder.block.cross_attn.num_heads: choice(6, 8, 12)