Use pruner in executor

not522 · Sep 17, 2020 · b6dfc4e · b6dfc4e
1 parent 0c4e5bd
commit b6dfc4e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/allennlp/allennlp_jsonnet.py b/allennlp/allennlp_jsonnet.py
@@ -45,7 +45,11 @@ def objective(trial):
     if version.parse(allennlp.__version__) < version.parse("1.0.0"):
         raise RuntimeError("AllenNLP>=1.0.0 is required for this example.")
 
-    study = optuna.create_study(direction="maximize")
+    study = optuna.create_study(
+        direction="maximize",
+        storage="sqlite:///allennlp.db",
+        pruner=optuna.pruners.HyperbandPruner(),
+    )
     study.optimize(objective, n_trials=50, timeout=600)
 
     print("Number of finished trials: ", len(study.trials))

diff --git a/allennlp/classifier.jsonnet b/allennlp/classifier.jsonnet
@@ -68,5 +68,10 @@ local ENCODER = CNN_FIELDS(
     },
     patience: 2,
     validation_metric: '+accuracy',
+    epoch_callbacks: [
+      {
+        type: 'optuna_pruner',
+      },
+    ],
   },
 }