3 Upvotes

Using Sklearn RFE to Select Features

Python
Data Preparation for Models

Recursive Feature Elimination (RFE) is a feature selection method that enables fitting a model based on a selected number of top ranking features from the training data. Sklearn works with any model type that assigns weights to features.

In the example below we use a Random Forest Regression algorithm. Once initialised, this gets passed to the RFE method along with the number of features we want to select, in this case 8. After this, like any other Sklearn we can fit the model and make predictions by calling fit and predict.

Lastly, we combine the feature names from our training data and the ranking array from RFE (rfe.ranking_) into a dataframe and print this to analyse how RFE has ranked the features.

from sklearn.ensemble import RandomForestRegressor
from sklearn.feature_selection import RFE

rf = RandomForestRegressor(random_state=101)

rfe = RFE(rf, n_features_to_select=8)
rfe = rfe.fit(X_train, y_train)

predictions = rfe.predict(X_test)

#Print feature rankings
feature_rankings = pd.DataFrame({'feature_names':np.array(X_train.columns),'feature_ranking':rfe.ranking_})
print(feature_rankings)

By detro - Last Updated Dec. 6, 2021, 7:54 p.m.

Did you find this snippet useful?

Sign up to bookmark this in your snippet library

COMMENTS
RELATED SNIPPETS
Scale Data Using Standard Scaler in Sklearn
Python
Data Preparation for Models

Sklearn

3
Create Dummy Variables with Pandas
Python
Data Preparation for Models

Pandas

1
Top Contributors
75