257 lines
7.2 KiB
Python
257 lines
7.2 KiB
Python
import os
|
|
import json
|
|
import timeit
|
|
import pickle
|
|
import itertools
|
|
from abc import ABC, abstractmethod
|
|
from pathlib import Path
|
|
from multiprocessing import cpu_count
|
|
|
|
import numpy as np
|
|
|
|
|
|
def get_from_config():
|
|
"""Get benchmarks configuration from the config.json file"""
|
|
current_path = Path(__file__).resolve().parent
|
|
|
|
config_path = current_path / "config.json"
|
|
with open(config_path, "r") as config_file:
|
|
config_file = "".join(line for line in config_file if line and "//" not in line)
|
|
config = json.loads(config_file)
|
|
|
|
profile = os.getenv("SKLBENCH_PROFILE", config["profile"])
|
|
|
|
n_jobs_vals_env = os.getenv("SKLBENCH_NJOBS")
|
|
if n_jobs_vals_env:
|
|
n_jobs_vals = eval(n_jobs_vals_env)
|
|
else:
|
|
n_jobs_vals = config["n_jobs_vals"]
|
|
if not n_jobs_vals:
|
|
n_jobs_vals = list(range(1, 1 + cpu_count()))
|
|
|
|
cache_path = current_path / "cache"
|
|
cache_path.mkdir(exist_ok=True)
|
|
(cache_path / "estimators").mkdir(exist_ok=True)
|
|
(cache_path / "tmp").mkdir(exist_ok=True)
|
|
|
|
save_estimators = os.getenv("SKLBENCH_SAVE_ESTIMATORS", config["save_estimators"])
|
|
save_dir = os.getenv("ASV_COMMIT", "new")[:8]
|
|
|
|
if save_estimators:
|
|
(cache_path / "estimators" / save_dir).mkdir(exist_ok=True)
|
|
|
|
base_commit = os.getenv("SKLBENCH_BASE_COMMIT", config["base_commit"])
|
|
|
|
bench_predict = os.getenv("SKLBENCH_PREDICT", config["bench_predict"])
|
|
bench_transform = os.getenv("SKLBENCH_TRANSFORM", config["bench_transform"])
|
|
|
|
return (
|
|
profile,
|
|
n_jobs_vals,
|
|
save_estimators,
|
|
save_dir,
|
|
base_commit,
|
|
bench_predict,
|
|
bench_transform,
|
|
)
|
|
|
|
|
|
def get_estimator_path(benchmark, directory, params, save=False):
|
|
"""Get path of pickled fitted estimator"""
|
|
path = Path(__file__).resolve().parent / "cache"
|
|
path = (path / "estimators" / directory) if save else (path / "tmp")
|
|
|
|
filename = (
|
|
benchmark.__class__.__name__
|
|
+ "_estimator_"
|
|
+ "_".join(list(map(str, params)))
|
|
+ ".pkl"
|
|
)
|
|
|
|
return path / filename
|
|
|
|
|
|
def clear_tmp():
|
|
"""Clean the tmp directory"""
|
|
path = Path(__file__).resolve().parent / "cache" / "tmp"
|
|
for child in path.iterdir():
|
|
child.unlink()
|
|
|
|
|
|
class Benchmark(ABC):
|
|
"""Abstract base class for all the benchmarks"""
|
|
|
|
timer = timeit.default_timer # wall time
|
|
processes = 1
|
|
timeout = 500
|
|
|
|
(
|
|
profile,
|
|
n_jobs_vals,
|
|
save_estimators,
|
|
save_dir,
|
|
base_commit,
|
|
bench_predict,
|
|
bench_transform,
|
|
) = get_from_config()
|
|
|
|
if profile == "fast":
|
|
warmup_time = 0
|
|
repeat = 1
|
|
number = 1
|
|
min_run_count = 1
|
|
data_size = "small"
|
|
elif profile == "regular":
|
|
warmup_time = 1
|
|
repeat = (3, 100, 30)
|
|
data_size = "small"
|
|
elif profile == "large_scale":
|
|
warmup_time = 1
|
|
repeat = 3
|
|
number = 1
|
|
data_size = "large"
|
|
|
|
@property
|
|
@abstractmethod
|
|
def params(self):
|
|
pass
|
|
|
|
|
|
class Estimator(ABC):
|
|
"""Abstract base class for all benchmarks of estimators"""
|
|
|
|
@abstractmethod
|
|
def make_data(self, params):
|
|
"""Return the dataset for a combination of parameters"""
|
|
# The datasets are cached using joblib.Memory so it's fast and can be
|
|
# called for each repeat
|
|
pass
|
|
|
|
@abstractmethod
|
|
def make_estimator(self, params):
|
|
"""Return an instance of the estimator for a combination of parameters"""
|
|
pass
|
|
|
|
def skip(self, params):
|
|
"""Return True if the benchmark should be skipped for these params"""
|
|
return False
|
|
|
|
def setup_cache(self):
|
|
"""Pickle a fitted estimator for all combinations of parameters"""
|
|
# This is run once per benchmark class.
|
|
|
|
clear_tmp()
|
|
|
|
param_grid = list(itertools.product(*self.params))
|
|
|
|
for params in param_grid:
|
|
if self.skip(params):
|
|
continue
|
|
|
|
estimator = self.make_estimator(params)
|
|
X, _, y, _ = self.make_data(params)
|
|
|
|
estimator.fit(X, y)
|
|
|
|
est_path = get_estimator_path(
|
|
self, Benchmark.save_dir, params, Benchmark.save_estimators
|
|
)
|
|
with est_path.open(mode="wb") as f:
|
|
pickle.dump(estimator, f)
|
|
|
|
def setup(self, *params):
|
|
"""Generate dataset and load the fitted estimator"""
|
|
# This is run once per combination of parameters and per repeat so we
|
|
# need to avoid doing expensive operations there.
|
|
|
|
if self.skip(params):
|
|
raise NotImplementedError
|
|
|
|
self.X, self.X_val, self.y, self.y_val = self.make_data(params)
|
|
|
|
est_path = get_estimator_path(
|
|
self, Benchmark.save_dir, params, Benchmark.save_estimators
|
|
)
|
|
with est_path.open(mode="rb") as f:
|
|
self.estimator = pickle.load(f)
|
|
|
|
self.make_scorers()
|
|
|
|
def time_fit(self, *args):
|
|
self.estimator.fit(self.X, self.y)
|
|
|
|
def peakmem_fit(self, *args):
|
|
self.estimator.fit(self.X, self.y)
|
|
|
|
def track_train_score(self, *args):
|
|
if hasattr(self.estimator, "predict"):
|
|
y_pred = self.estimator.predict(self.X)
|
|
else:
|
|
y_pred = None
|
|
return float(self.train_scorer(self.y, y_pred))
|
|
|
|
def track_test_score(self, *args):
|
|
if hasattr(self.estimator, "predict"):
|
|
y_val_pred = self.estimator.predict(self.X_val)
|
|
else:
|
|
y_val_pred = None
|
|
return float(self.test_scorer(self.y_val, y_val_pred))
|
|
|
|
|
|
class Predictor(ABC):
|
|
"""Abstract base class for benchmarks of estimators implementing predict"""
|
|
|
|
if Benchmark.bench_predict:
|
|
|
|
def time_predict(self, *args):
|
|
self.estimator.predict(self.X)
|
|
|
|
def peakmem_predict(self, *args):
|
|
self.estimator.predict(self.X)
|
|
|
|
if Benchmark.base_commit is not None:
|
|
|
|
def track_same_prediction(self, *args):
|
|
est_path = get_estimator_path(self, Benchmark.base_commit, args, True)
|
|
with est_path.open(mode="rb") as f:
|
|
estimator_base = pickle.load(f)
|
|
|
|
y_val_pred_base = estimator_base.predict(self.X_val)
|
|
y_val_pred = self.estimator.predict(self.X_val)
|
|
|
|
return np.allclose(y_val_pred_base, y_val_pred)
|
|
|
|
@property
|
|
@abstractmethod
|
|
def params(self):
|
|
pass
|
|
|
|
|
|
class Transformer(ABC):
|
|
"""Abstract base class for benchmarks of estimators implementing transform"""
|
|
|
|
if Benchmark.bench_transform:
|
|
|
|
def time_transform(self, *args):
|
|
self.estimator.transform(self.X)
|
|
|
|
def peakmem_transform(self, *args):
|
|
self.estimator.transform(self.X)
|
|
|
|
if Benchmark.base_commit is not None:
|
|
|
|
def track_same_transform(self, *args):
|
|
est_path = get_estimator_path(self, Benchmark.base_commit, args, True)
|
|
with est_path.open(mode="rb") as f:
|
|
estimator_base = pickle.load(f)
|
|
|
|
X_val_t_base = estimator_base.transform(self.X_val)
|
|
X_val_t = self.estimator.transform(self.X_val)
|
|
|
|
return np.allclose(X_val_t_base, X_val_t)
|
|
|
|
@property
|
|
@abstractmethod
|
|
def params(self):
|
|
pass
|