El modelo KNN que usa pyfunc devuelve ModuleNotFoundError o FileNotFoundError

Problema

Ha creado un modelo de Sklearn mediante KNeighborsClassifier y está usando pyfunc para ejecutar una predicción.

Por ejemplo:

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()

La predicción devuelve un mensaje de error ModuleNotFoundError: No module named 'sklearn.neighbors._classification'.

La predicción también puede devolver un mensaje de error FileNotFoundError: [Errno 2] No usable temporary directory found.

Causa

Cuando se registra un modelo de KNN, todos los puntos de datos usados para el entrenamiento se guardan como parte del archivo pickle.

Si el modelo se entrena con millones de registros, todos los datos se agregan al modelo, lo que puede aumentar considerablemente su tamaño. Un modelo entrenado con millones de registros puede suman fácilmente varios GB.

pyfunc intenta cargar todo el modelo en la memoria caché del ejecutor al ejecutar una predicción.

Si el modelo es demasiado grande para caber en la memoria, aparece uno de los mensajes de error anteriores.

Solución

Debe usar un algoritmo basado en árboles, como Bosque aleatorio o XGBoost para reducir el tamaño de los datos en un modelo de KNN.

Si tiene datos desequilibrados, intente un método de muestreo como SMOTE al entrenar un algoritmo basado en árboles.