Guardar los modelos de Scikit-learn en disco e importarlo en otra sesión

Una vez entrenado un modelo de aprendizaje automático con Scikit-learn puede surgir la necesidad de guardar este para usar en otra sesión. Posiblemente durante el proceso ha sido necesario cargar los datos, seleccionar las carteristas más relevantes, ajustar los hiperparámetros y comparar varios algoritmos de aprendizaje. Algo que no querremos repetir cada vez que necesitemos obtener una predicción de este modelo. Siendo la opción ideal poder guardar los modelos de Scikit-learn en archivos para su posterior recuperación, esto es, crear una versión persistente del modelo. Lo que se puede obtener fácilmente gracias a la librería joblib.

Creación de un modelo de ejemplo

Antes de poder guardar un modelo en un archivo es necesario crear uno. Para ello se puede crear un conjunto de datos aleatorio con la función make_regression() y entrenado con un modelo de regresión cómo puede ser Lasso. Lo que se puede obtener con el siguiente código.

from sklearn.datasets import make_regression
from sklearn.linear_model import Lasso

X, y = make_regression(n_features=3, random_state=0)

model = Lasso().fit(X, y)

model.predict([[1,1,1]])

array([106.07180807])

Un ejemplo básico en el que se usan los datos directamente, sin ningún procesado, para la creación del modelo. Obteniendo una predicción de 106,07 para el vector [1, 1, 1].

Guardar el modelo de Scikit-learn en disco

Ahora que tenemos un modelo en nuestra sesión de Python es posible guardar este en un archivo con la función dump() de joblib. Una función que únicamente necesita dos parámetros, en primer lugar, el propio modelo y en segundo el nombre del archivo donde se guardará. Así, para guardar el modelo solamente se tendrá que escribir y lanzar el siguiente código.

from joblib import dump

dump(model, 'regression.joblib')

Una vez lanzado este, aparecerá en la carpeta de trabajo un nuevo archivo llamado regression.joblib con el modelo.

Importar el modelo desde disco

Para importar el modelo se puede recurrir a la función load() de joblib. Función que requiere como parámetro el nombre del archivo y devuelve el objeto guardado, en nuestro caso el modelo. Para lo que se puede ejecutar el siguiente código.

from joblib import load

regression_model = load('regression.joblib')

regression_model.predict([[1,1,1]])

array([106.07180807])

Ejemplo en el que, además de importar el modelo, también se ha obtenido la predicción del modelo para el vector [1, 1, 1], obteniendo el mismo resultado que en la sesión original.

Compresión del modelo

En el ejemplo que hemos visto el modelo no es complejo, por lo que el archivo resultante solamente ocupa unos 711 bytes, pero en casos más complejos puede ser interesante comprimir el archivo para ahorrar espacio. Algo que se puede conseguir mediante el uso del parámetro compress de la función dump(). Un parámetro que por defecto tiene el valor 0, sin compresión, y puede tener cualquier valor hasta 9, máxima compresión. Así el siguiente código genera un archivo de 490 bytes con el mismo modelo, esto es, se ha reducido el tamaño en un 31%.

dump(model, 'compress.joblib', compress=9)

Mediante el parámetro compress también se puede indicar el tipo de compresor utilizado ('zlib', 'gzip', 'bz2', ...). Para lo que habrá que pasar una tupla dónde el primer elemento sea un texto con el tipo de compresión y el segundo un valor entre 0 y 9 con el nivel de compresión.

Conclusiones

En esta entrada se ha visto cómo se puede guardar los modelos de Scikit-learn en disco con joblib para posteriormente importarlos en otra sesión. Algo que puede ser útil cuando necesitamos usar estos modelos en otras máquinas para obtener predicciones.

Imagen de Alexas from Pixabay