我有两个colors
和excluded_colors
的数据。
colors
包含所有颜色excluded_colors
包含一些我希望从我的训练集中排除的颜色。
我试图将数据分成训练集和测试集,并确保excluded_colors
中的颜色不在我的训练集中,但存在于测试集中。
为了实现上述目标,我做了如下操作
var colors = spark.sql("""
select colors.*
from colors
LEFT JOIN excluded_colors
ON excluded_colors.color_id = colors.color_id
where excluded_colors.color_id IS NULL
"""
)
val trainer: (Int => Int) = (arg:Int) => 0
val sqlTrainer = udf(trainer)
val tester: (Int => Int) = (arg:Int) => 1
val sqlTester = udf(tester)
val rsplit = colors.randomSplit(Array(0.7, 0.3))
val train_colors = splits(0).select("color_id").withColumn("test",sqlTrainer(col("color_id")))
val test_colors = splits(1).select("color_id").withColumn("test",sqlTester(col("color_id")))
然而,我意识到,通过上述excluded_colors
中的颜色完全被忽略了。它们甚至不在我的测试集中。
我如何将数据分成70/30,同时确保excluded_colors
中的颜色不是在训练中而是在测试中存在。
我们要做的是从训练集中删除"被排除的颜色",但将它们放在测试中,并将训练/测试分割为70/30。
我们需要的是一点数学。
给定总数据集(TD)和排除的颜色数据集(E),我们可以说对于训练数据集(Tr)和测试数据集(Ts):
|Tr| = x * (|TD|-|E|)
|Ts| = |E| + (1-x) * |TD|
我们知道|Tr| = 0.7 |TD|
因此x = 0.7 |TD| / (|TD| - |E|)
现在我们知道了采样因子x
,我们可以说:
Tr = (TD-E).sample(withReplacement = false, fraction = x)
// where (TD - E) is the result of the SQL expr above
Ts = TD.sample(withReplacement = false, fraction = 0.3)
// we sample the test set from the original dataset