From a886e9d94ff332cdc2c6d9774488f8b4b25d8656 Mon Sep 17 00:00:00 2001
From: dong <2938690875@qq.com>
Date: Fri, 1 Nov 2024 00:17:46 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20function.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 function.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/function.py b/function.py
index 9d768a6..8c1ba8b 100644
--- a/function.py
+++ b/function.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import train_test_split
 
 
 id = "202210423130"
@@ -89,13 +90,12 @@ def preprocess_data(X):
 
 
 if __name__ == "__main__":
-    # 加载训练集
-    X_train, y_train = load_data("./data/horse-colic.data")
-    X_train = preprocess_data(X_train)
+    # 加载数据集
+    X, y = load_data("./data/horse-colic.data")
+    X = preprocess_data(X)
 
-    # 加载测试集
-    X_test, y_test = load_data("./data/horse-colic.test")
-    X_test = preprocess_data(X_test)
+    # 使用 train_test_split 打乱并随机抽取数据
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
 
     # 训练对率回归模型
     weights = train_logreg(X_train, y_train, learning_rate=0.1, iterations=1000, lambda_reg=0.01)