更新 function.py

function.py

@@ -1,14 +1,11 @@
 import os
 import csv
 import numpy as np
-import torch
-import torch.nn as nn
-import torch.optim as optim
+import tensorflow as tf
 from sklearn.model_selection import train_test_split
 from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, roc_auc_score
 from sklearn.preprocessing import StandardScaler
 from imblearn.over_sampling import SMOTE
-from torch.utils.data import DataLoader, TensorDataset
 
 # 读取 CSV 数据的函数
 def load_data(file_path):
@@ -37,15 +34,6 @@ def evaluate_model(y_true, y_pred, y_prob=None):
         auc = roc_auc_score(y_true, y_prob)
         print(f"AUC Score: {auc:.4f}")
 
-# 定义逻辑回归模型
-class LogisticRegressionModel(nn.Module):
-    def __init__(self, input_dim):
-        super(LogisticRegressionModel, self).__init__()
-        self.linear = nn.Linear(input_dim, 1)  # 线性层
-
-    def forward(self, x):
-        return torch.sigmoid(self.linear(x))  # 使用 Sigmoid 激活函数
-
 # 主函数
 def main():
     # 加载数据
@@ -64,37 +52,21 @@ def main():
     X_train_res = scaler.fit_transform(X_train_res)
     X_test = scaler.transform(X_test)
 
-    # 将数据转换为 Tensor
-    X_train_tensor = torch.FloatTensor(X_train_res)
-    y_train_tensor = torch.FloatTensor(y_train_res).view(-1, 1)
-    X_test_tensor = torch.FloatTensor(X_test)
-    y_test_tensor = torch.FloatTensor(y_test).view(-1, 1)
+    # 创建 Keras 模型
+    model = tf.keras.Sequential([
+        tf.keras.layers.InputLayer(input_shape=(X_train_res.shape[1],)),
+        tf.keras.layers.Dense(1, activation='sigmoid')  # 逻辑回归输出
+    ])
 
-    # 创建数据加载器
-    train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
-    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-
-    # 初始化模型、损失函数和优化器
-    input_dim = X_train_res.shape[1]
-    model = LogisticRegressionModel(input_dim)
-    criterion = nn.BCELoss()  # 二元交叉熵损失
-    optimizer = optim.SGD(model.parameters(), lr=0.01)
+    # 编译模型
+    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
 
     # 训练模型
-    model.train()
-    for epoch in range(100):  # 训练 100 个 epoch
-        for inputs, labels in train_loader:
-            optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            loss.backward()
-            optimizer.step()
+    model.fit(X_train_res, y_train_res, epochs=100, batch_size=32, verbose=1)
 
     # 评估模型
-    model.eval()
-    with torch.no_grad():
-        y_pred_prob = model(X_test_tensor).numpy()
-        y_pred = (y_pred_prob > 0.5).astype(int)
+    y_pred_prob = model.predict(X_test).flatten()  # 获取预测概率
+    y_pred = (y_pred_prob > 0.5).astype(int)
 
     # 评估模型
     evaluate_model(y_test, y_pred, y_pred_prob)