機器學習筆記——資料集分割

在模型訓練之前，要首先劃分訓練集與測試集，如何對原始資料集進行訓練集與測試集的劃分？訓練集與測試集的比例各佔多少？如何保證各自內部標籤分佈平衡都會影響模型訓練的最終效果。

好在R和Python中有現成的資料集分割函式，避免手動寫函式導致劃分比例不合理、訓練集與測試集的樣本的結構與總體不均衡的問題。

R語言中caTools包中的sample.split函式可以用來自動將原始資料集分割成訓練集和測試集。

方法一——caTools中的sample.split函式

library("caTools")
set.seed(123)

data(iris)
table(iris$Species)

setosa versicolor virginica
50 50 50

split = sample.split(iris$Species,SplitRatio = .8)
train_data = subset(iris,split == TRUE)
test_data = subset(iris,split == FALSE)

table(train_data$Species)
setosa versicolor virginica
40 40 40

table(test_data$Species)
setosa versicolor virginica
10 10 10

劃分方法二——createDataPartition函式

library("caret")
split1 <- createDataPartition(y=iris$Species,p=0.8,list = FALSE)

train_data <- iris[split1,]
table(train_data$Species)
 setosa versicolor virginica 
 40 40 40 


test_data <- iris[-split1,]
table(test_data$Species)
 setosa versicolor virginica 
 10 10 10

可以看到無論是caTools包中的sample.split函式還是caret包中的createDataPartition函式，都針對分類標籤做了混合後的分層隨機抽樣，這樣可以保證訓練集與測試集內的各類標籤分佈比例與樣本總體的分佈比例嚴格一致，否則如果僅僅使用sample函式，無法達到分層隨機抽樣的目的。

Python的sk-learn庫中也有現成的資料集分割工具可用。

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_iris
import pandas as pd

iris = load_iris()
data = iris[`data`]
iris_data = pd.DataFrame(
 data = data,
 columns = [`sepal_length`,`sepal_width`,`petal_length`,`petal_width`]
 )
iris_data["Species"] = iris[ `target`]
iris_data["Species"] = iris_data["Species"].map({0:"setosa",1:"versicolor",2:"virginica"})

iris_data["Species"].value_counts()
versicolor 50
virginica 50
setosa 50
Name: Species, dtype: int64

x,y = iris_data.iloc[:,0:-1],iris_data.iloc[:,-1]
train_data,test_data,train_target,test_target = train_test_split(x,y,test_size = 0.2,stratify = y)

train_target.value_counts()
test_target.value_counts()

這裡有一個細節，其中train_test_split函式中有兩個特殊的引數：shuffle和stratify。

shuffle引數相當於對原始資料進行混合抽樣，相當於撲克牌發牌前的洗牌，保證隨機性。

stratify引數則可以保證訓練集&測試集中樣本標籤結構比例與指定的總體中樣本標籤結構比例一致，特別是在原始資料中樣本標籤分佈不均衡時非常有用，達到分層隨機抽樣的目的。

而對於訓練集與測試集的劃分比例上來看，目前沒有一個統一的比例標準，在機器學習類的參考資料中，推薦的比例是訓練集和測試集的比例保持在1:2~1:4之間為佳。

倘若使用驗證集，則建議保持6:2:2左右的訓練集、驗證集及測試集比例。

原文釋出時間為：2018-09-23

本文作者：杜雨

本文來自雲棲社群合作伙伴“資料小魔方”，瞭解相關資訊可以關注“資料小魔方”。