使用 pyfunc 的 KNN 模型返回 ModuleNotFoundErrorFileNotFoundError

问题

你使用 KNeighborsClassifier 创建了一个 Sklearn 模型,并使用 pyfunc 运行预测。

例如:

import mlflow.pyfunc
pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, result_type='string')
predicted_df = merge.withColumn("prediction", pyfunc_udf(*merge.columns[1:]))
predicted_df.collect()

预测返回一条 ModuleNotFoundError: No module named 'sklearn.neighbors._classification' 错误消息。

预测也可能返回 FileNotFoundError: [Errno 2] No usable temporary directory found 错误消息。

原因

记录 KNN 模型后,用于训练的所有数据点都将保存为 pickle 文件的一部分。

如果用数百万条记录对模型进行训练,所有这些数据都会添加到模型中,这将极大地增加其大小。 经过数百万条记录训练的模型可以轻松地合计多个 GB。

pyfunc 尝试在运行预测时将整个模型加载到执行程序的缓存中。

如果模型过大,内存容纳不下,则会出现上述错误消息之一。

解决方案

应使用基于树的算法,如随机林或 XGBoost,减少对 KNN 模型数据的采样。

如果数据不平衡,请在训练基于树的算法时,尝试使用 SMOTE 之类的采样方法。