Como salvar seu modelo de Machine Learning

Se você precisa rodar seu Jupyter notebook inteiro a cada vez que precisa fazer uma nova previsão, saiba que tem uma maneira fácil de salvar o seu modelo de machine learning e usá-lo diretamente a partir de um arquivo.

Notebooks são a melhor maneira que eu conheço para estruturar projetos de Data Science. Eles aceitam markdown, html e células de código – permitem documentação, organização e que sejam replicáveis por qualquer pessoa.

No entanto, se a cada vez que você precisa usar aquele algoritmo XGBoost treinado você ainda roda todas as células novamente, saiba que existe uma biblioteca chamada pickle capaz de exportar estruturas de dados para um arquivo e importá-las em qualquer outra máquina.

Ganhe tempo e aprenda a usar (em poucas linhas) esse pacote nos seus trabalhos de Data Science.

Salvando modelos de machine learning com Pickle

Na verdade, não são apenas modelos de machine learning que podem ser salvos em um arquivo no seu disco. Qualquer objeto em Python (strings, dicionários, listas, arrays) pode ser facilmente exportado.

Opickle permite a serialização de objetos, ou seja, os transforma em sequências de bytes. Armazenando toda informação necessária, consegue reconstruir objetos aplicando uma lógica interna.

Para salvar um modelo, usa-se o método dump() – basta informar a variável referente ao objeto e o nome do arquivo a ser criado, como no exemplo abaixo.

# salvar modelo
import pickle
# salvar o modelo XGBoost (xgb_model) no arquivo sale_xgboost.pkl
with open('sale_xgboost.pkl', 'wb') as file:
    pickle.dump(xgb_model, file)

Pronto. Um arquivo com o nome sale_xgboost.pkl foi criado na mesma pasta onde o código foi executado. Basta salvar ele, jogar para deploy ou usar quando precisar.

Carregando modelos de machine learning com Pickle

O mais comum em machine learning é você treinar e avaliar o modelo até ele estar satisfatório, e jogar para deploy dentro de uma aplicação web, por exemplo. Para isso, precisamos carregar o modelo.

Esse caminho também é muito direto. Usando o método load() você consegue pegar o objeto que foi serializado (transformado em uma sequência de bytes) e convertê-lo novamente para o tipo original

# Carregar modelo
with open('sale_xgboost.pkl', 'rb') as f:
    model = pickle.load(f)

Já que o modelo estava treinado e validado, uma vez carregado ele já esta pronto para ser usado. No caso da célula acima, o modelo de machine learning foi atribuído à variável model.

Para se fazer uma previsão, é só usar o método o model.predict(X_test) (por exemplo) como você faria normalmente.

Além dos modelos de machine learning

Como eu te disse, você consegue salvar qualquer objeto do seu código Python, e não apenas modelos e algoritmos treinados em projetos de Ciência de Dados.

Às vezes, tem-se apenas o resultado de uma query ou um array do numpy que você deseja salvar para reaproveitar depois (afinal, quando você reiniciar a IDE as variáveis serão eliminadas).

Para todos esses casos, o pacote pickle vai te ajudar muito. O fato de ele vir junto com o próprio Python já ajuda muito, pois nem se perde tempo instalando componentes a mais

Espero que tenha gostado da dica e que comece a usar essa técnica a partir de agora 🙂