如何使用scikit-learn的train_test_split()函式分割資料集


在機器學習和資料科學的廣闊領域中,我們會遇到一些看似微不足道但實際上至關重要的任務。其中一項重要任務就是將資料分割成訓練集和驗證集——這是建立有效預測模型的基礎步驟。Scikit-learn是一個著名的Python機器學習庫,它擁有一個功能強大的函式`train_test_split()`,可以輕鬆完成此任務。本文旨在指導您使用scikit-learn的`train_test_split()`函式分割資料。

語法

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
  • X和y分別代表屬性矩陣和目標向量。

  • test_size表示預留給驗證集的初始資料比例(通常為0.2或20%)。

  • random_state初始化控制資料分割的內部隨機數生成器。

使用scikit-learn的train_test_split()函式分割資料

此操作簡化了將資料分割成訓練集和驗證集的任務。以下是它的工作原理

  • 此示例演示了一個基本的訓練測試分割,測試大小為20%。

  • 80%的資料將構成訓練集(X_train和y_train),其餘20%將構成測試集(X_test和y_test)。

  • 每個集合中的精確資料點取決於輸入資料和隨機狀態。

a. 基本訓練測試分割

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

輸出

X_train, X_test, y_train, and y_test

b. 分層訓練測試分割

  • 此示例展示了一個分層訓練測試分割。

  • stratify引數確保訓練集和測試集中的每個類的比例與原始資料集中的每個類的比例相同。

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, stratify=y)

輸出

X_train, X_test, y_train, and y_test

c. 訓練-驗證-測試分割

  • 此示例演示了一個訓練-驗證-測試分割。資料首先被分成訓練集(60%的資料)和臨時集(40%的資料)。

  • 然後,臨時集進一步被分成驗證集和測試集,每個集包含原始資料的20%。

X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.4)
X_validation, X_test, y_validation, y_test = train_test_split(X_temp, y_temp, test_size=0.5)

輸出

X_train, X_validation, X_test, y_train, y_validation, and y_test

d. 帶洗牌的分割

  • 此示例顯示了一個啟用了洗牌的訓練測試分割。shuffle引數確保在分割槽之前對資料進行隨機洗牌。

  • 每個集合中的資料點將被隨機洗牌。

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=True)

輸出

X_train, X_test, y_train, and y_test

e. 使用特定資料集的分割

  • 此示例展示了一個具有特定隨機狀態的訓練測試分割。random_state引數設定隨機數生成器的種子,確保每次執行程式碼時生成相同的訓練測試分割。

  • 由於隨機狀態固定,每個集合中的資料點在多次執行中將保持一致。

   X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

輸出

X_train, X_test, y_train, and y_test

結論

來自scikit-learn的`train_test_split()`函式簡化了將資料分割成訓練集和驗證集的任務。這是一個強大的函式,擁有許多引數,可以根據任務需要定製分割槽。

`train_test_split()` 的靈活性使其能夠適應不同的環境,使其成為任何資料科學家或機器學習從業者必備的工具。透過掌握如何熟練地使用此函式,

更新於:2023年8月28日

瀏覽量:146

啟動您的職業生涯

透過完成課程獲得認證

開始學習
廣告