You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
25 lines
645 B
Python
25 lines
645 B
Python
3 years ago
|
from turtle import pd
|
||
|
|
import joblib as jl
|
||
|
|
|
||
|
|
from sklearn.datasets import load_iris
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
from sklearn.ensemble import RandomForestClassifier
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
# Use iris dataset
|
||
|
|
iris = load_iris()
|
||
|
|
X, y = iris.data, iris.target
|
||
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y)
|
||
|
|
clr = RandomForestClassifier()
|
||
|
|
# Fit
|
||
|
|
clr.fit(X_train, y_train)
|
||
|
|
# Serialize the classifier to pickle file
|
||
|
|
jl.dump(clr, "./output/model.pkl", compress=9)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
print("Building iris model...")
|
||
|
|
main()
|
||
|
|
print("Model trained and dumped as pickle file.")
|