Motivação
Quem nunca salvou um modelo do scikit-learn em formato pickle que atire a primeira pedra. Piadas a parte, sei de muitos utilizam pickle como formato padrão para persistência de modelos de machine learning (ML).
Porém eu estava assistindo a seguinte talk do Nubank
quando engenheiros da área de MLOps de lá solta a seguinte frase:
Utilizamos pickle e, infelizmente, ainda não conseguimos nos desvencilhar
Essa foi uma afirmação que me deu um certo choque e fez com que eu me perguntasse: se não for em pickle, como que se faz então? Esse é o pivô que motivou esse artigo, porém, ainda existiram duas outras situações.
Um dia normal no trabalho
Um belo dia uma colega de trabalho me chamou perguntando se eu tinha algum material legal para aprender um pouco mais sobre backend e construção de APIs, mas que também abordasse aspectos do nosso trabalho que também envolvia infraestrutura. Sugeri ela de fazer a rinha de backend. Foi nesse momento que também comentei que eu estudava Go e que tinha submetido uma solução pra rinha usando Go — muito motivado pela economia de recursos, dada as restrições da rinha.
Passam alguns dias, uma semana talvez, e a empresa começa uma grande frente de corte de custos com computação em nuvem visando atingir a meta de lucro proposta para o final daquele ano. Foi aí que essa colega me pergunta: porque a gente não faz o deploy de todos os modelos que temos aqui em Go? Uma vez que Go utiliza muito menos recurso que Python, a gente consegue fazer uma grande economia.
Na ocasião disse que não seria possível, porque os modelos são escritos em Python, salvos em pickle e o pickle é um formato binário Python, ou seja, sem possibilidade de ser aberto em código Go. Não só por esse impeditivo, mas, também pelo esforço para migrarmos todo o CI/CD já construído, dentre muitos outros detalhes que deveriam ser alterados para viabilizar. Por outro lado, confesso que fiquei curioso para saber se isso seria possível ou não.
IA integrada no celular
A pouco tempo comecei a ver algumas propagandas de IA diretamente integrada em celulares e fiquei me perguntando como que fazem o deploy localmente no celular, visto que geralmente os modelos são treinados em Python e muitas vezes precisam de mais recursos que um celular consegue oferecer. Em um primeiro momento achei que por debaixo eles chamavam o runtime do Python pra poder rodar os modelos, mas isso iria criar um overhead grande, principalmente em dispositivos mobile, que possuem restrição de recursos.
Foi aí que eu topei com o formato ONNX....
O que é ONNX?
ONNX significa Open Neural Network Exchange. Diferente do pickle, o ONNX é um formato aberto para representação de modelos de ML e IA na forma de um grafo. Isso é possível dado que no formato são definidos operadores comuns para modelos de ML e IA. O ONNX tem ganhado destaque por sua capacidade de promover a interoperabilidade entre diferentes frameworks. Por exemplo, codar uma rede neural em PyTorch e abrir utilizando Tensorflow. Para se aprofundar nos conceitos do ONNX, recomendo visitar a documentação oficial.
Além disso, ainda existe o ONNX runtime que é mantido pela Microsoft nesse repositório. Através desse runtime, é possível rodar um modelo ONNX em qualquer ambiente no qual seja possível instalar o runtime. Dessa forma, tornando possível fazer o treinamento de um modelo em ambiente nuvem utilizando o PyTorch e para o deploy desse modelo em um ambiente mobile bastaria utilizar o runtime do ONNX para dispositivos mobile. Isso já me responde como devem fazer o deploy de IAs diretamente em dispositivos mobile.
Mas, ainda faltou a resposta para a primeira pergunta....
É possível abrir modelos Python em Golang?
Se depender somente do ONNX, sim, é possível. Mas, existem algumas etapas que precisam ser cumpridas antes. Nesse artigo irei expor um exemplo bem básico de como atingir esse objetivo, considerando uma regressão logística treinada em cima do iris dataset.
Treinamento do modelo
Como o iris dataset é um dataset que já vem muito bem tratado, não vou ficar me alongando quanto a isso. Vou direto para o treinamento do modelo. Mas, caso tenha interesse em ver todo o processo, é só acessar o link do repositório aqui.
O modelo foi criado através de um pipeline do próprio scikit-learn e possui 3 etapas:
Um
SimpleImputerpara lidar com valores vazios. Dessa forma, se uma das features tiver um valor vazio, ele será completado com a média dos outros dois valores;Adição de um
StandardScalerpara normalização dos dados;O estimador, ou modelo, de fato que fará a predição. Para esse exemplo, foi escolhido uma regressão logística, dada a simplicidade e agilidade no treinamento do modelo.
pipeline = Pipeline(steps = [
('Imputer', SimpleImputer(strategy='mean', keep_empty_features=True)),
('normalization', StandardScaler()),
('estimator', LogisticRegression())
]
)De forma a treinar o modelo com os melhores parâmetros, também utilizei um GridSearchCV antes de chamar o .fit .
parameters = {
'estimator__solver': ['newton-cg'],
'estimator__tol': [ 0.0001, 0.003, 0.01],
'estimator__penalty': [None, 'l2'],
}
model = GridSearchCV(estimator=pipeline,
param_grid=parameters,
scoring= {"AUC": "roc_auc_ovr"},
refit="AUC",
cv=5,
verbose=1,
error_score='raise')
model = model.fit(X_train, y_train)Fazendo um rápido teste para ver se o modelo está performando bem.
y_pred = model.predict(X_test)
test_acc_score = accuracy_score(y_test, y_pred)
logging.info(f"Accuracy test score: {test_acc_score}")
## INFO:root:Accuracy test score: 0.9666666666666667Bom, 96% de acurácia é bem razoável. Dessa forma, conseguimos concluir que o modelo conseguiu aprender com os exemplos. Podem me julgar por utilizar acurácia para avaliar o modelo e mais nada. Eu sei que pode ter sofrido com overfit ou qualquer outro tipo de viés. Mas, aqui o foco é no ONNX e não na qualidade do modelo em si. Lembre-se disso.
Convertendo o modelo para ONNX
Antes de mais nada, esse foi o modelo resultante do nosso treinamento e é ele que queremos salvar no formato ONNX.
GridSearchCV(cv=5, error_score='raise',
estimator=Pipeline(steps=[('Imputer',
SimpleImputer(keep_empty_features=True)),
('normalization', StandardScaler()),
('estimator', LogisticRegression())]),
param_grid={'estimator__penalty': [None, 'l2'],
'estimator__solver': ['newton-cg'],
'estimator__tol': [0.0001, 0.003, 0.01]},
refit='AUC', scoring={'AUC': 'roc_auc_ovr'}, verbose=1)Para salvar o modelo scikit-learn em ONNX podemos utilizar uma biblioteca com um nome bem sugestivo: skl2onnx . Como utilizamos um GridSearchCV para fazer o treinamento do nosso modelo com cross validation, precisamos acessar o melhor estimador.
Outro ponto importante é criar o tipo do tensor que será o nosso input. Lembre-se que o ONNX é uma representação na forma de um grafo do nosso modelo. Então, esse tensor de input tem dimensão [1,4] , basicamente é um vetor (os matemáticos que me julguem). Dessa forma, nosso tensor de input tem 4 posições, sendo uma para cada feature na qual o modelo foi treinado e terá como input para as predições.
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
initial_type = [('float_input', FloatTensorType([None, 4]))]
onnx_model = convert_sklearn(model.best_estimator_, initial_types=initial_type)O tensor de output é gerado automaticamente, mas, basicamente será um tensor com dois outros dentro. O primeiro contendo a classe prevista pelo modelo, já o segundo com a distribuição de probabilidades para cada uma das classes possíveis.
Já dando um pequeno spoiler de como abrir modelos ONNX usando Python, podemos fazer uma predição com o modelo convertido para ONNX e comparar com o nosso modelo em pickle para vermos se pelo menos os números batem e garantir que nossa conversão foi bem sucedida.
import onnxruntime as rt
sess = rt.InferenceSession(f"logistic_regression_iris.onnx")
input_name = sess.get_inputs()[0].name
print("Input Name:", input_name)
## Input Name: float_input ## -> O nome bate com o que informamos no snippet anterior
X_test_np = X_test.astype(np.float32).to_numpy()
predictions = sess.run(None, {input_name: X_test_np})
print(predictions)
## [array([2, ...], dtype=int64), [{0: 4.3682564864866436e-05, 1: 0.18593932688236237, 2: 0.8140169978141785}...]]
pred = model.predict(X_test_np)
print(pred)
## [2. ...]Uma vez convertido, podemos utilizar o Neutron para termos uma visualização do grafo gerado.
Grafo que representa as etapas executadas pelo modelo convertido para ONNX
Abrindo o modelo usando Go
Existem algumas bibliotecas que auxiliam na abertura de modelos ONNX em Go através do ONNX runtime, a que eu usei nesse exemplo está disponível nesse repositório. Mas, como o ONNX runtime é uma biblioteca C, é necessário utilizar uma versão do Go que tenha suporte ao CGo, assim como informar no código Go o caminho onde o binário dessa biblioteca se encontra na máquina. Dessa forma, a biblioteca que eu utilizei, basicamente, é um wrapper do ONNX runtime em C para Go.
import ort "github.com/yalue/onnxruntime_go"
func InitializeEnvironment(shared_lib_path string) {
ort.SetSharedLibraryPath(shared_lib_path)
err := ort.InitializeEnvironment()
if err != nil {
panic(err)
}
}
func CreateSession(model_path string) (*ort.DynamicAdvancedSession, error) {
options := ort.SessionOptions{}
defer options.Destroy()
sess, err := ort.NewDynamicAdvancedSession(model_path, []string{"float_input"}, []string{"output_label", "output_probability"}, &options)
if err != nil {
return nil, err
}
return sess, nil
}Porém, o que deu mais trabalho foi criar algumas funções para facilitar a predição utilizando o modelo. Abaixo você encontra dois exemplos, o primeiro que retorna somente a classe prevista pelo modelo, o segundo retorna a distribuição de probabilidades.
var MAPPED_LABELS = map[int64]string{
0: "iris-setosa",
1: "iris-versicolor",
2: "iris-virginica",
}
func Predict(input_values []float32, sess *ort.DynamicAdvancedSession, input_shape ...int64) (string, error) {
input_tensor, err := Create_input_tensor(input_values, input_shape...)
if err != nil {
return "", err
}
defer input_tensor.Destroy()
output_values := []ort.Value{nil, nil}
err = sess.Run([]ort.Value{input_tensor}, output_values)
if err != nil {
return "", err
}
defer output_values[0].Destroy()
defer output_values[1].Destroy()
label_tensor := output_values[0].(*ort.Tensor[int64])
defer label_tensor.Destroy()
return MAPPED_LABELS[label_tensor.GetData()[0]], nil
}func Predict_proba(input_values []float32, sess *ort.DynamicAdvancedSession, input_shape ...int64) (map[string]float32, error) {
input_tensor, err := Create_input_tensor(input_values, input_shape...)
if err != nil {
return map[string]float32{}, err
}
defer input_tensor.Destroy()
output_values := []ort.Value{nil, nil}
err = sess.Run([]ort.Value{input_tensor}, output_values)
if err != nil {
return map[string]float32{}, err
}
defer output_values[0].Destroy()
defer output_values[1].Destroy()
sequence := output_values[1].(*ort.Sequence)
probability_maps, err := sequence.GetValues()
if err != nil {
panic(err)
}
probabilities := make(map[string]float32)
for i := range probability_maps {
m := probability_maps[i].(*ort.Map)
keys, values, err := m.GetKeysAndValues()
if err != nil {
panic(err)
}
keys_tensor := keys.(*ort.Tensor[int64])
values_tensor := values.(*ort.Tensor[float32])
for j, key := range keys_tensor.GetData() {
value := values_tensor.GetData()[j]
probabilities[MAPPED_LABELS[key]] = value
}
}
return probabilities, nil
}Deployment
A fim de testar qual seria a melhor forma de fazer o deployment desse modelo, entenda melhor como aguentar mais carga consumindo a menor quantidade de recurso possível, resolvi criar 3 APIs. Sendo:
A primeira utilizando FastAPI e servindo inferências utilizando o modelo em pickle. De forma "tradicional";
A segunda também em FastAPI, porém, servindo o modelo em ONNX; e, por fim
A terceira servindo o modelo em ONNX utilizando Go e o go-chi como router;
Com as APIs prontas, optei por fazer o deploy delas localmente utilizando um arquivo docker-compose.yaml restringindo os recursos das três para:
0.5 CPUs; e
128Mb de memória.
services:
python-api-pkl:
container_name: poc-python-api-pkl
build:
dockerfile: ./python-api-pkl/Dockerfile
context: .
environment:
WEBSERVER_PORT: 5000
ports:
- 5001:5000
deploy:
resources:
limits:
cpus: '0.50'
memory: 128M
python-api-onnx:
container_name: poc-python-api-onnx
build:
dockerfile: ./python-api-onnx/Dockerfile
context: .
environment:
WEBSERVER_PORT: 5000
ports:
- 5002:5000
deploy:
resources:
limits:
cpus: '0.50'
memory: 128M
go-api-onnx:
container_name: poc-go-api-onnx
build:
dockerfile: ./go-api-onnx/build/Dockerfile
context: .
environment:
WEBSERVER_PORT: 5000
SHARED_LIB_PATH: /go-api-onnx/onnxruntime-linux-x64-1.20.1/lib/libonnxruntime.so
MODEL_PATH_ONNX: /models/logistic_regression_iris.onnx
OTEL_SERVICE_NAME: poc-go-api-onnx
ports:
- 5003:5000
deploy:
resources:
limits:
cpus: '0.50'
memory: 128MTeste de carga
A fim de gerar a carga nas APIs, utilizei o K6. Essa escolha se dá pela simplicidade de criação de testes de carga, além de ser muito leve de rodar. Dado que toda a infraestrutura está dentro da minha própria máquina e os recursos são bem limitados.
A carga configurada para o K6 foi a seguinte:
Começar com 100 requisições por minuto durante 2 minutos; depois incrementar linearmente a carga para;
4000 requisições por minuto durante 2 minutos;
8000 requisições por minuto durante 2 minutos;
16000 requisições por minuto durante 2 minutos.
Dessa forma, o teste tem uma duração total de 8 minutos. Como payload, utilizei um número aleatório para cada feature, porém, garantindo que esse valor esteja dentro do intervalo mínimo e máximo do dataset. Um ponto a se ressaltar é que eu desconsiderei qualquer validação do retorno das APIs durante o teste de carga.
import http from 'k6/http';
import { sleep } from 'k6';
function randomFloat(min, max) {
return Math.random() * (max - min) + min;
}
export let options = {
scenarios: {
load_increase: {
executor: 'ramping-arrival-rate',
startRate: 100,
timeUnit: '1m',
preAllocatedVUs: 50,
maxVUs: 10000,
stages: [
{ duration: '2m', target: 4000 },
{ duration: '2m', target: 8000 },
{ duration: '2m', target: 16000 },
],
startTime: '2m',
}
}
};
export default function () {
const urls = [
'http://localhost:5001/predict',
'http://localhost:5002/predict',
'http://localhost:5003/predict'
];
urls.forEach(url => {
const payload = JSON.stringify({
sepal_length: randomFloat(4.3, 7.9),
sepal_width: randomFloat(2, 4.4),
petal_length: randomFloat(1, 6.9),
petal_width: randomFloat(0.1, 2.5),
});
const params = {
headers: { 'Content-Type': 'application/json' }
};
http.post(url, payload, params);
});
sleep(1);
}Coleta de dados e telemetria
A fim de conseguir monitorar as aplicações, todas elas foram instrumentadas utilizando opentelemetry. As métricas escolhidas para medir a performance foram:
Latência no p95;
Requisições por minuto (RPM);
Consumo de memória; e
Consumo de CPU.
As APIs exportam as métricas geradas para um Opentelemetry collector. Sendo o collector responsável pela conversão do formato OTLP exportado pelas APIs para o formato do Prometheus. Por sua vez, o Prometheus faz o scrape das métricas das APIs através do collector e das métricas de containers expostas através de um container do cAdvisor. Por fim, para visualizar os dados, existe um container do Grafana que faz as queries no Prometheus e gera os dashboards.
Arquitetura da solução
Resultado do teste de carga
Resultados do teste de carga
Python pickle
A API feito em Python com o modelo no formato pickle foi a que mais consumiu memória e CPU, assim como também foi a que teve a maior elevação da latência no p95 conforme a carga aumentava. Um ponto importante de se levantar, é que com 128Mb de memória a API sempre ficou acima dos 90% de uso. Também podemos ver no gráfico de RPM que a curva de requisições lidadas por essa API foi a menor das três, porém, não muito abaixo das demais.
Python ONNX
Confesso que essa me surpreendeu. Não só pela latência no p95 ter sido bastante próxima da API em Go, mas por ter uma redução considerável no uso de memória e CPU em comparação com a API usando o modelo em pickle. Dessa forma, já se mostrando uma boa alternativa para se utilizar caso seja necessário uma redução dos custos de computação.
Go ONNX
Como era de se esperar, a API em Go foi a que teve o menor consumo de recursos computacionais em comparação com as outras duas APIs e manteve a latência no p95 estável do início ao fim do teste, um comportamento parecido com a API Python ONNX. Porém, o que me surpreendeu foi que ela não teve o melhor RPM. Ficando, inclusive, próximo da API Python pickle.
Conclusão
De forma bem objetiva, é possível abrir modelos treinados em Python em outras linguagens, sendo uma das formas utilizando o ONNX. Pode exigir um pouco mais do time tanto de ciência de dados quanto de MLOps para lidar com as particularidades desse formato. Porém, dada a economia de recursos que o formato ONNX pode promover, pode ser uma boa pedida.
Por outro lado, o teste que eu fiz foi somente utilizando um dos modelos mais simples que existem, uma regressão logística. Dessa forma, modelos mais complexos podem não ser suportados ou necessitar de etapas na conversão e carregamento do modelo que não cobri nesse artigo. Outro ponto é que a economia de recurso observada nos testes desse artigo podem não ser replicados para outros modelos.
Caso você tenha alguma sugestão de melhoria, sinta-se a vontade para entrar em contato ou abrir uma issue no repositório onde está o código para todas as APIs e o teste de carga.