models.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. # models imports
  2. from sklearn.model_selection import GridSearchCV
  3. from sklearn.linear_model import LogisticRegression
  4. from sklearn.ensemble import RandomForestClassifier, VotingClassifier
  5. from sklearn.neighbors import KNeighborsClassifier
  6. from sklearn.ensemble import GradientBoostingClassifier
  7. import sklearn.svm as svm
  8. def _get_best_model(X_train, y_train):
  9. Cs = [0.001, 0.01, 0.1, 1, 10, 100, 1000]
  10. gammas = [0.001, 0.01, 0.1, 1, 5, 10, 100]
  11. param_grid = {'kernel':['rbf'], 'C': Cs, 'gamma' : gammas}
  12. svc = svm.SVC(probability=True)
  13. clf = GridSearchCV(svc, param_grid, cv=10, scoring='accuracy', verbose=0)
  14. clf.fit(X_train, y_train)
  15. model = clf.best_estimator_
  16. return model
  17. def svm_model(X_train, y_train):
  18. return _get_best_model(X_train, y_train)
  19. def ensemble_model(X_train, y_train):
  20. svm_model = _get_best_model(X_train, y_train)
  21. lr_model = LogisticRegression(solver='liblinear', multi_class='ovr', random_state=1)
  22. rf_model = RandomForestClassifier(n_estimators=100, random_state=1)
  23. ensemble_model = VotingClassifier(estimators=[
  24. ('svm', svm_model), ('lr', lr_model), ('rf', rf_model)], voting='soft', weights=[1,1,1])
  25. ensemble_model.fit(X_train, y_train)
  26. return ensemble_model
  27. def ensemble_model_v2(X_train, y_train):
  28. svm_model = _get_best_model(X_train, y_train)
  29. knc_model = KNeighborsClassifier(n_neighbors=2)
  30. gbc_model = GradientBoostingClassifier(n_estimators=100, learning_rate=1.0, max_depth=1, random_state=0)
  31. lr_model = LogisticRegression(solver='liblinear', multi_class='ovr', random_state=1)
  32. rf_model = RandomForestClassifier(n_estimators=100, random_state=1)
  33. ensemble_model = VotingClassifier(estimators=[
  34. ('lr', lr_model),
  35. ('knc', knc_model),
  36. ('gbc', gbc_model),
  37. ('svm', svm_model),
  38. ('rf', rf_model)],
  39. voting='soft', weights=[1, 1, 1, 1, 1])
  40. ensemble_model.fit(X_train, y_train)
  41. return ensemble_model
  42. def get_trained_model(choice, X_train, y_train):
  43. if choice == 'svm_model':
  44. return svm_model(X_train, y_train)
  45. if choice == 'ensemble_model':
  46. return ensemble_model(X_train, y_train)
  47. if choice == 'ensemble_model_v2':
  48. return ensemble_model_v2(X_train, y_train)