add --no-parallel-split argument to higg benchmark

ogrisel · Dec 19, 2018 · e84cdc3 · e84cdc3
1 parent b466047
commit e84cdc3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/benchmarks/bench_higgs_boson.py b/benchmarks/bench_higgs_boson.py
@@ -21,6 +21,7 @@
 parser.add_argument('--learning-rate', type=float, default=1.)
 parser.add_argument('--subsample', type=int, default=None)
 parser.add_argument('--max-bins', type=int, default=255)
+parser.add_argument('--no-parallel-split', action="store_true", default=False)
 args = parser.parse_args()
 
 HERE = os.path.dirname(__file__)
@@ -85,7 +86,8 @@ def load_data():
                                          max_leaf_nodes=n_leaf_nodes,
                                          n_iter_no_change=None,
                                          random_state=0,
-                                         verbose=1, parallel_splitting=False)
+                                         verbose=1,
+                                         parallel_splitting=not args.no_parallel_split)
 pygbm_model.fit(data_train, target_train)
 toc = time()
 predicted_test = pygbm_model.predict(data_test)