models.py 2.2 KB

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