from abc import ABC import matplotlib.pyplot as plt from lightgbm import LGBMClassifier from sklearn import svm from sklearn.ensemble import RandomForestClassifier, VotingClassifier from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay, accuracy_score class ModelTrainer(ABC): """ 模型训练类 """ def __init__(self, toForecast) -> None: super().__init__() self._toForecast = toForecast ################### # Step 3: Train the model def test_model(self, strMethod, classifier, X_test, y_test): print(strMethod + " ====== test results ======") y_pred = classifier.predict(X_test) result0 = confusion_matrix(y_test, y_pred, labels=[0, 1]) print(strMethod + " Confusion Matrix:") print(result0) result1 = classification_report(y_test, y_pred, zero_division=1.0) print(strMethod + " Classification Report:") print(result1) result2 = accuracy_score(y_test, y_pred) print(strMethod + " Accuracy:", result2) cm_display = ConfusionMatrixDisplay(confusion_matrix=result0, display_labels=['Down', 'Up']) cm_display.plot() plt.title(strMethod + ' Accuracy: ' + f'{result2:.0%}') plt.show() def train_random_forest(self, X_train, y_train, X_test, y_test): classifier = RandomForestClassifier() classifier.fit(X_train, y_train) if not self._toForecast: self.test_model('Random Forest', classifier, X_test, y_test) return classifier def train_GBT(self, X_train, y_train, X_test, y_test): # Gradient Boosted Tree classifierGBT = LGBMClassifier() classifierGBT.fit(X_train, y_train) if not self._toForecast: self.test_model('Gradient Boosted Tree', classifierGBT, X_test, y_test) return classifierGBT def train_SVC(self, X_train, y_train, X_test, y_test): # Support Vector Machines classifierSVC = svm.SVC() classifierSVC.fit(X_train, y_train) if not self._toForecast: self.test_model('Support Vector Machines', classifierSVC, X_test, y_test) return classifierSVC def ensemble_model(self, rf_model, gbt_model, svc_model, X_train, y_train, X_test, y_test): # Create a dictionary of our models estimators = [('rf', rf_model), ('gbt', gbt_model), ('svc', svc_model)] # Create our voting classifier, inputting our models ensemble = VotingClassifier(estimators, voting='hard') # fit model to training data ensemble.fit(X_train, y_train) if not self._toForecast: self.test_model('Ensemble Model', ensemble, X_test, y_test) return ensemble