如何划分机器学习的训练集和测试集

使用机器学习算法时，通常需要把数据分为训练集和测试集，本文介绍R语言的三种实现方法，并通过示例进行学习。

使用R内置方法

依据sample函数生成指定概率的true和false的向量，然后利用该向量过滤数据集得到训练集和测试集，语法如下：

# 设置随机种子，使得示例可以重复
set.seed(1)

# df是要分割的数据集
# 使用 70% 数据集作为训练集，30% 作为测试集
sample <- sample(c(TRUE, FALSE), nrow(df), replace=TRUE, prob=c(0.7,0.3))
train  <- df[sample, ]
test   <- df[!sample, ]

下面示例把iris数据分为训练集和测试集：

# 设置随机种子，使得示例可以重复
set.seed(1)

data(iris)

# 使用 70% 数据集作为训练集，30% 作为测试集
sample <- sample(c(TRUE, FALSE), nrow(iris), replace=TRUE, prob=c(0.7,0.3))
train  <- iris[sample, ]
test   <- iris[!sample, ]

dim(iris)
# 150   5

dim(train)
# 106   5

dim(test)
# 44  5

从输出可以看到，总共150条记录，训练集大概包括 106 / 150 = 70.6% 。

使用caTools包

使用caTools包提供了sample.split函数，可以轻松进行数据分离。
subset函数返回满足条件的向量、矩阵或数据帧的子集，subset(x, subset, ...) 其中subset参数指定过滤条件：

library(caTools)

# 设置随机种子，使得示例可以重复
set.seed(1)

# df是要分割的数据集
# 使用 70% 数据集作为训练集，30% 作为测试集
sample <- sample.split(df$any_column, SplitRatio = 0.7)
train  <- subset(df, sample == TRUE)
test   <- subset(df, sample == FALSE)

我们用这种方法对iris数据集进行划分：

# 设置随机种子，使得示例可以重复
data(iris)
set.seed(2)

# df是要分割的数据集
# 使用 70% 数据集作为训练集，30% 作为测试集
sample <- sample.split(iris$Species, SplitRatio = 0.7)

train  <- subset(iris, sample == TRUE)
test   <- subset(iris, sample == FALSE)

dim(iris)
# 150   5

dim(train)
# 105   5

dim(test)
# 45  5

运行结果与上面示例差不多，需要说明的是sample.split函数，其语法如下：sample.split( Y, SplitRatio = 2/3, group = NULL )

其中Y为x向量类型，因此需要随意指定数据集的一列生成向量，SplitRatio有默认值。

使用dplyr包

下面我们来看看如何使用dplyr包实现同样功能。

library(dplyr)

# 设置随机种子，使得示例可以重复
set.seed(1)

df <- iris

# 依据数据集总行数生成序号
df$id <- 1:nrow(df)

# 使用 70% 数据集作为训练集，30% 作为测试集
train <- df %>% dplyr::slice_sample(prop = 0.70)

# 反向连接生成测试机
test  <- dplyr::anti_join(df, train, by = 'id')

dim(iris)
# 150   5

dim(train)
# 105   5

dim(test)
# 45   5

生成结果与上面一致，我们说明下slice_sample()函数，用于随机选择行，语法如下：

slice_sample(.data, ..., n, prop, weight_by = NULL, replace = FALSE)

• prop 用于指定概率
• n 指定数量
• replace 是否放回抽样，默认不放回