DOC accelerate plot_gradient_boosting_regularization.py example #21598 (#21611)

sply88 · web-flow · commit f19bf4c406e9 · 2021-11-29T13:44:51.000+01:00
* accelerate plot_gradient_boosting_regularization.py example #21598 * speed up by less samples and less trees * use train_test_split instead of slicing
diff --git a/examples/ensemble/plot_gradient_boosting_regularization.py b/examples/ensemble/plot_gradient_boosting_regularization.py
@@ -31,18 +31,17 @@
 from sklearn import ensemble
 from sklearn import datasets
 
+from sklearn.model_selection import train_test_split
 
-X, y = datasets.make_hastie_10_2(n_samples=12000, random_state=1)
-X = X.astype(np.float32)
+X, y = datasets.make_hastie_10_2(n_samples=4000, random_state=1)
 
 # map labels from {-1, 1} to {0, 1}
 labels, y = np.unique(y, return_inverse=True)
 
-X_train, X_test = X[:2000], X[2000:]
-y_train, y_test = y[:2000], y[2000:]
+X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.8, random_state=0)
 
 original_params = {
-    "n_estimators": 1000,
+    "n_estimators": 400,
     "max_leaf_nodes": 4,
     "max_depth": None,
     "random_state": 2,
@@ -53,17 +52,17 @@
 
 for label, color, setting in [
     ("No shrinkage", "orange", {"learning_rate": 1.0, "subsample": 1.0}),
-    ("learning_rate=0.1", "turquoise", {"learning_rate": 0.1, "subsample": 1.0}),
+    ("learning_rate=0.2", "turquoise", {"learning_rate": 0.2, "subsample": 1.0}),
     ("subsample=0.5", "blue", {"learning_rate": 1.0, "subsample": 0.5}),
     (
-        "learning_rate=0.1, subsample=0.5",
+        "learning_rate=0.2, subsample=0.5",
         "gray",
-        {"learning_rate": 0.1, "subsample": 0.5},
+        {"learning_rate": 0.2, "subsample": 0.5},
     ),
     (
-        "learning_rate=0.1, max_features=2",
+        "learning_rate=0.2, max_features=2",
         "magenta",
-        {"learning_rate": 0.1, "max_features": 2},
+        {"learning_rate": 0.2, "max_features": 2},
     ),
 ]:
     params = dict(original_params)
