-
Notifications
You must be signed in to change notification settings - Fork 0
/
auto_fl.py
40 lines (30 loc) · 1 KB
/
auto_fl.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from flaml.automl import AutoML
from pmlb import fetch_data
def main():
df = fetch_data("adult")
X = df.drop(["target"], axis=1)
print(X.shape)
y = df["target"]
print(y.shape)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=21
)
automl = AutoML()
automl.fit(X_train=X_train, y_train=y_train, task="classification")
# retrieve best config and best learner
print("Best ML leaner:", automl.best_estimator)
print("Best hyperparmeter config:", automl.best_config)
# get predictions
preds = automl.predict(X_test)
# print evaluation scores
print(classification_report(y_test, preds))
# Visualize feature importance
plt.barh(
automl.model.estimator.feature_name_,
automl.model.estimator.feature_importances_,
)
if __name__ == "__main__":
main()