Merge branch 'main' into model-parsing

Author: pavelzw

.github/workflows/benchmark.yml

@@ -26,17 +26,18 @@ jobs:
           cache-name: benchmark-model-cache
         with:
           path: examples/benchmark_models
-          key: models-${{ hashFiles('benchmark.py') }}
+          key: models-${{ hashFiles('benchmark.py', '.github/workflows/benchmark.yml') }}
       - name: Run benchmark
         shell: bash -el {0}
         run: |
-          echo "_(benchmark **${{ github.run_id }}** / attempt **${{ github.run_attempt }}**)_" >> benchmark.md
-          python benchmark.py >> benchmark.md
+          python benchmark.py
+          echo "_(benchmark **${{ github.run_id }}** / attempt **${{ github.run_attempt }}**)_" >> result.md
+          cat benchmark.md >> result.md
       - name: Comment on PR
         if: github.event_name == 'pull_request'
         uses: thollander/actions-comment-pull-request@v2
         with:
-          filePath: benchmark.md
+          filePath: result.md
           comment_tag: benchmark
       - name: Add benchmark to Job Summary
         run: cat benchmark.md >> "$GITHUB_STEP_SUMMARY"

.github/workflows/build.yml

@@ -31,13 +31,13 @@ jobs:
           name: artifact
           path: dist
       - name: Publish package on TestPyPi
-        uses: pypa/gh-action-pypi-publish@a3a3bafbb3e5a75a854ae1bc53ae128cf22c4af4
+        uses: pypa/gh-action-pypi-publish@48b317d84d5f59668bb13be49d1697e36b3ad009
         with:
           user: __token__
           password: ${{ secrets.TEST_PYPI_TOKEN }}
           repository_url: https://test.pypi.org/legacy/
       - name: Publish package on PyPi
-        uses: pypa/gh-action-pypi-publish@a3a3bafbb3e5a75a854ae1bc53ae128cf22c4af4
+        uses: pypa/gh-action-pypi-publish@48b317d84d5f59668bb13be49d1697e36b3ad009
         with:
           user: __token__
           password: ${{ secrets.PYPI_TOKEN }}

.gitignore

@@ -1,3 +1,5 @@
+benchmark.md
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

benchmark.py

@@ -1,4 +1,5 @@
 import io
+import itertools
 import lzma
 import pickle
 import textwrap
@@ -11,6 +12,7 @@
 
 from examples.utils import generate_dataset
 from slim_trees.lgbm_booster import dump_lgbm
+from slim_trees.pickling import get_pickled_size
 from slim_trees.sklearn_tree import dump_sklearn
 
 MODELS_PATH = "examples/benchmark_models"
@@ -20,51 +22,93 @@ def load_model(model_name: str, generate: Callable) -> Any:
     model_path = Path(f"{MODELS_PATH}/{model_name}.pkl")
 
     if model_path.exists():
+        print(f"Loading model `{model_name}.pkl` from disk...")
         with open(model_path, "rb") as f:
             return pickle.load(f)
 
+    print(f"Training model `{model_name}`...")
     regressor = generate()
-    regressor.fit(*generate_dataset(n_samples=10000))
+    regressor.fit(
+        *generate_dataset(n_samples=10000),
+    )
+    size = get_pickled_size(regressor, "no", pickle.dump)
+    print(f"Trained model {model_name}. Size {size / 2**20:.2f} MB")
+
     model_path.parent.mkdir(parents=True, exist_ok=True)
+
     with open(model_path, "wb") as f:
         pickle.dump(regressor, f)
     return regressor
 
 
-def train_gb_sklearn() -> GradientBoostingRegressor:
+def train_sklearn_rf_20m() -> RandomForestRegressor:
+    return load_model(
+        "sklearn_rf_20m",
+        lambda: RandomForestRegressor(
+            n_estimators=100, max_leaf_nodes=1700, random_state=42, n_jobs=-1
+        ),
+    )
+
+
+def train_sklearn_rf_200m() -> RandomForestRegressor:
+    return load_model(
+        "sklearn_rf_200m",
+        lambda: RandomForestRegressor(n_estimators=275, random_state=42, n_jobs=-1),
+    )
+
+
+def train_sklearn_rf_1g() -> RandomForestRegressor:
     return load_model(
-        "gb_sklearn",
-        lambda: GradientBoostingRegressor(n_estimators=2000, random_state=42),
+        "sklearn_rf_1g",
+        lambda: RandomForestRegressor(
+            n_estimators=1500, max_leaf_nodes=10000, random_state=42, n_jobs=-1
+        ),
+    )
+
+
+def train_sklearn_gb_2m() -> GradientBoostingRegressor:
+    return load_model(
+        "sklearn_gb_2m",
+        lambda: GradientBoostingRegressor(
+            n_estimators=2000, random_state=42, verbose=True
+        ),
     )
 
 
-def train_model_sklearn() -> RandomForestRegressor:
+def train_lgbm_gbdt_2m() -> lgb.LGBMRegressor:
     return load_model(
-        "rf_sklearn",
-        lambda: RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1),
+        "lgbm_gbdt_2m", lambda: lgb.LGBMRegressor(n_estimators=1000, random_state=42)
     )
 
 
-def train_gbdt_lgbm() -> lgb.LGBMRegressor:
+def train_lgbm_gbdt_5m() -> lgb.LGBMRegressor:
     return load_model(
-        "gbdt_lgbm", lambda: lgb.LGBMRegressor(n_estimators=2000, random_state=42)
+        "lgbm_gbdt_5m",
+        lambda: lgb.LGBMRegressor(n_estimators=2000, random_state=42),
     )
 
 
-def train_gbdt_large_lgbm() -> lgb.LGBMRegressor:
+def train_lgbm_gbdt_20m() -> lgb.LGBMRegressor:
     return load_model(
-        "gbdt_large_lgbm",
-        lambda: lgb.LGBMRegressor(n_estimators=20000, random_state=42),
+        "lgbm_gbdt_20m",
+        lambda: lgb.LGBMRegressor(n_estimators=8000, random_state=42),
     )
 
 
-def train_rf_lgbm() -> lgb.LGBMRegressor:
+def train_lgbm_gbdt_100m() -> lgb.LGBMRegressor:
     return load_model(
-        "rg_lgbm",
+        "lgbm_gbdt_100m",
+        lambda: lgb.LGBMRegressor(n_estimators=35000, random_state=42),
+    )
+
+
+def train_lgbm_rf_10m() -> lgb.LGBMRegressor:
+    return load_model(
+        "lgbm_rf_10m",
         lambda: lgb.LGBMRegressor(
             boosting_type="rf",
-            n_estimators=100,
-            num_leaves=1000,
+            n_estimators=700,
+            num_leaves=8000,
             random_state=42,
             bagging_freq=5,
             bagging_fraction=0.5,
@@ -74,12 +118,9 @@ def train_rf_lgbm() -> lgb.LGBMRegressor:
 
 
 def benchmark(func: Callable, *args, **kwargs) -> float:
-    times = []
-    for _ in range(5):
-        start = time.perf_counter()
-        func(*args, **kwargs)
-        times.append(time.perf_counter() - start)
-    return min(times)
+    start = time.perf_counter()
+    func(*args, **kwargs)
+    return time.perf_counter() - start
 
 
 def benchmark_model(  # noqa: PLR0913
@@ -99,11 +140,13 @@ def benchmark_model(  # noqa: PLR0913
 
     model = train_func()
 
+    print(f"Benchmarking naive implementation of `{name}`...")
     naive_dump_time = benchmark(base_dumps_func, model)
     naive_pickled = base_dumps_func(model)
     naive_pickled_size = len(naive_pickled)
     naive_load_time = benchmark(base_loads_func, naive_pickled)
 
+    print(f"Benchmarking our implementation of `{name}`...")
     our_dump_time = benchmark(dumps_func, model)
     our_pickled = dumps_func(model)
     our_pickled_size = len(our_pickled)
@@ -210,18 +253,39 @@ def loads_lzma(data):
         dumps_lzma,
         loads_lzma,
     )
-    models_to_benchmark = [
-        ("sklearn rf", train_model_sklearn) + dumps_sklearn_args,
-        ("sklearn rf LZMA", train_model_sklearn) + dumps_sklearn_lzma_args,
-        ("sklearn gb", train_gb_sklearn) + dumps_sklearn_args,
-        ("sklearn gb LZMA", train_gb_sklearn) + dumps_sklearn_lzma_args,
-        ("LGBM gbdt", train_gbdt_lgbm) + dumps_lgbm_args,
-        ("LGBM gbdt LZMA", train_gbdt_lgbm) + dumps_lgbm_lzma_args,
-        ("LGBM gbdt large", train_gbdt_large_lgbm) + dumps_lgbm_args,
-        ("LGBM gbdt large LZMA", train_gbdt_large_lgbm) + dumps_lgbm_lzma_args,
-        ("LGBM rf", train_rf_lgbm) + dumps_lgbm_args,
-        ("LGBM rf LZMA", train_rf_lgbm) + dumps_lgbm_lzma_args,
+    models = [
+        ("sklearn rf 20M", train_sklearn_rf_20m),
+        ("sklearn rf 200M", train_sklearn_rf_200m),
+        ("sklearn rf 1G", train_sklearn_rf_1g),
+        ("sklearn gb 2M", train_sklearn_gb_2m),
+        ("lgbm gbdt 2M", train_lgbm_gbdt_2m),
+        ("lgbm gbdt 5M", train_lgbm_gbdt_5m),
+        ("lgbm gbdt 20M", train_lgbm_gbdt_20m),
+        ("lgbm gbdt 100M", train_lgbm_gbdt_100m),
+        ("lgbm rf 10M", train_lgbm_rf_10m),
     ]
+
+    def get_dumps_args(model_name, train_func):
+        if "sklearn" in model_name:
+            return (model_name, train_func) + dumps_sklearn_args
+        elif "lgbm" in model_name:
+            return (model_name, train_func) + dumps_lgbm_args
+        else:
+            raise ValueError(f"Unknown model name: {model_name}")
+
+    def get_dumps_args_lzma(model_name, train_func):
+        if "sklearn" in model_name:
+            return (model_name + " lzma", train_func) + dumps_sklearn_lzma_args
+        elif "lgbm" in model_name:
+            return (model_name + " lzma", train_func) + dumps_lgbm_lzma_args
+        else:
+            raise ValueError(f"Unknown model name: {model_name}")
+
+    models_to_benchmark = itertools.chain.from_iterable(
+        [[get_dumps_args(*model), get_dumps_args_lzma(*model)] for model in models]
+    )
     benchmark_results = [benchmark_model(*args) for args in models_to_benchmark]
-    print("Base results / Our results / Change")
-    print(format_benchmarks_results_table(benchmark_results))
+    results_str = format_benchmarks_results_table(benchmark_results)
+    with open("benchmark.md", "w") as f:
+        f.write("Base results / Our results / Change\n")
+        f.write(results_str)