Spaces:
Runtime error
Runtime error
File size: 1,638 Bytes
3d6943b | 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 41 42 43 44 45 46 47 48 49 50 51 52 53 | import optuna
import mlflow
import mlflow.sklearn
import numpy as np
from src.utils.mlflow_utils import setup_mlflow
from sklearn.ensemble import RandomForestRegressor
from sklearn.pipeline import Pipeline
from sklearn.metrics import mean_squared_error
def tune_random_forest(X_train, X_test, y_train, y_test, preprocessor):
setup_mlflow("AQI_Prediction")
def objective(trial):
with mlflow.start_run(nested=True):
params = {
"n_estimators": trial.suggest_int("n_estimators", 50, 300),
"max_depth": trial.suggest_int("max_depth", 5, 30),
"min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
"min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 5),
"random_state": 42,
"n_jobs": -1
}
mlflow.log_params(params)
model = RandomForestRegressor(**params)
pipeline = Pipeline([
("preprocessor", preprocessor),
("model", model)
])
pipeline.fit(X_train, y_train)
y_pred = pipeline.predict(X_test)
rmse = mean_squared_error(y_test, y_pred) ** 0.5
mlflow.log_metric("rmse", rmse)
return rmse
with mlflow.start_run(run_name="Optuna_RF_Search"):
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=50)
mlflow.log_params(study.best_params)
mlflow.log_metric("best_rmse", study.best_value)
return study.best_params
|