在pyspark中保存并加载两个ML模型

时间:2017-08-01 16:18:49

标签: python apache-spark pyspark apache-spark-ml

首先,我创建两个ML算法并将它们保存到两个单独的文件中。请注意,两个模型都基于相同的数据帧。 feature_1feature_2是从同一数据集中提取的不同功能集。

import sys
from pyspark.ml.classification import RandomForestClassificationModel

trainer_1 = RandomForestClassifier(featuresCol="features_1")
trainer_2 = RandomForestClassifier(featuresCol="features_2")
model_1 = trainer_1.fit(df_training_data)
model_2 = trainer_2.fit(df_training_data)

model_1.save(sys.argv[1])
model_2.save(sys.argv[2])

然后,当我后来想要使用模型时,我必须从它们各自的路径加载它们,提供路径f.ex.通过sys.argv。

import sys
from pyspark.ml.classification import RandomForestClassificationModel

model_1 = RandomForestClassificationModel.load(sys.argv[1])
model_2 = RandomForestClassificationModel.load(sys.argv[2])

我想要的是一种优雅的方式,能够将这两个模型一起保存在同一条路径中。我希望这主要是为了让用户每次保存和加载时都不必跟踪两个单独的路径名。这两个模型紧密相连,通常会一起创建并一起使用,因此它们只是一种模型。

这是管道的用途吗?

1 个答案:

答案 0 :(得分:5)

我想出了一种方法,只需将它们放在一个文件夹中即可。然后用户只需要提供并知道该文件夹的路径。

import sys
import os
from pyspark.ml.classification import RandomForestClassificationModel

trainer_1 = RandomForestClassifier(featuresCol="features_1")
trainer_2 = RandomForestClassifier(featuresCol="features_2")
model_1 = trainer_1.fit(df_training_data)
model_2 = trainer_2.fit(df_training_data)

path = 'model_rfc'
os.mkdir(path)
model_1.save(os.path.join(sys.argv[1], 'model_1'))
model_2.save(os.path.join(sys.argv[1], 'model_2'))

名称model_1model_2是硬编码的,用户无需知道。

import sys
import os
from pyspark.ml.classification import RandomForestClassificationModel

model_1 = RandomForestClassificationModel.load(os.path.join(sys.argv[1], 'model_1'))
model_2 = RandomForestClassificationModel.load(os.path.join(sys.argv[1], 'model_2'))

这应该可以解决问题。这是最好的方法吗?还是有更好的方法可以使用Spark库的功能将模型捆绑在一起?