02 - Caso prático usando Python para treinar um modelo simples (Iris) e exportar para ONNX

Continuação do artigo:
01 - Por que python é tão usado para I.A.

Exemplo prático

Vou montar um exemplo completo, na pegada “copiar, ajustar e rodar”:

Python para treinar um modelo simples (Iris)
Exportar para ONNX
API em Python (FastAPI + ONNX Runtime) servindo o modelo
Tudo orquestrado com Docker + docker-compose

1. Estrutura de pastas do projeto

ml-onnx-example/
  ml/
    train.py
    requirements.txt
  api/
    app.py
    requirements.txt
  model/              # será preenchida pelo treino (model.onnx)
  Dockerfile.ml
  Dockerfile.api
  docker-compose.yml

2. Script de treino em Python (gera `model.onnx`)

`ml/requirements.txt`

torch
scikit-learn

`ml/train.py`

import os
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn import datasets
import numpy as np

# Modelo simples para o dataset Iris
class IrisNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(4, 16)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(16, 3)  # 3 classes

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x  # logits

def main():
    # Carrega dataset Iris
    iris = datasets.load_iris()
    X = iris.data.astype(np.float32)   # shape (150, 4)
    y = iris.target.astype(np.int64)   # shape (150,)

    X_tensor = torch.from_numpy(X)
    y_tensor = torch.from_numpy(y)

    model = IrisNet()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)

    model.train()
    epochs = 200

    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_tensor)
        loss = criterion(outputs, y_tensor)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 50 == 0:
            print(f"Epoch [{epoch+1}/{epochs}] - Loss: {loss.item():.4f}")

    # Garante que pasta model exista
    os.makedirs("model", exist_ok=True)

    # Exporta para ONNX
    model.eval()
    dummy_input = torch.randn(1, 4)  # 4 features do Iris
    onnx_path = os.path.join("model", "model.onnx")

    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input": {0: "batch_size"},
            "output": {0: "batch_size"}
        },
        opset_version=17
    )

    print(f"Modelo ONNX salvo em: {onnx_path}")

if __name__ == "__main__":
    main()

Esse script:

Treina uma rede neural pequena no Iris.
Exporta um modelo ONNX com:
- input: tensor [batch_size, 4] (float32)
- output: tensor [batch_size, 3] (logits para 3 classes)

3. API para servir o modelo ONNX (FastAPI + ONNX Runtime)

`api/requirements.txt`

fastapi
uvicorn[standard]
onnxruntime
numpy

`api/app.py`

from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import onnxruntime as ort

app = FastAPI(
    title="Iris ONNX API",
    version="1.0.0",
    description="API de exemplo servindo modelo Iris em ONNX"
)

# Carrega o modelo ONNX na inicialização da API
# O arquivo será montado em /app/model/model.onnx via docker-compose
SESSION = ort.InferenceSession(
    "model/model.onnx",
    providers=["CPUExecutionProvider"]
)

# nomes das classes só para ficar mais amigável
CLASS_NAMES = ["setosa", "versicolor", "virginica"]

class PredictRequest(BaseModel):
    features: list[float]  # 4 floats do Iris

class PredictResponse(BaseModel):
    class_index: int
    class_name: str
    probabilities: list[float]

def softmax(x: np.ndarray) -> np.ndarray:
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()

@app.get("/healthz")
def healthz():
    return {"status": "ok"}

@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest):
    # Garante que há 4 features
    if len(req.features) != 4:
        return {
            "class_index": -1,
            "class_name": "invalid_input",
            "probabilities": []
        }

    # Prepara entrada: shape (1, 4)
    x = np.array([req.features], dtype=np.float32)

    input_name = SESSION.get_inputs()[0].name  # "input"
    outputs = SESSION.run(None, {input_name: x})
    logits = outputs[0][0]  # shape (3,)

    probs = softmax(logits)
    class_idx = int(np.argmax(probs))
    class_name = CLASS_NAMES[class_idx]

    return PredictResponse(
        class_index=class_idx,
        class_name=class_name,
        probabilities=probs.tolist()
    )

Exemplo de requisição:

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"features": [5.1, 3.5, 1.4, 0.2]}'

4. Dockerfiles

`Dockerfile.ml` (container de treino)

FROM python:3.12-slim

WORKDIR /app

# Dependências
COPY ml/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Script de treino
COPY ml/train.py .

# Pasta onde o modelo será salvo (montada por volume)
RUN mkdir -p model

CMD ["python", "train.py"]

`Dockerfile.api` (container da API)

FROM python:3.12-slim

WORKDIR /app

# Dependências da API
COPY api/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Código da API
COPY api/app.py .

# O arquivo model/model.onnx será montado via volume
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

5. docker-compose: orquestração completa

`docker-compose.yml`

version: "3.9"

services:
  trainer:
    build:
      context: .
      dockerfile: Dockerfile.ml
    volumes:
      - ./model:/app/model
    # você roda esse serviço manualmente para treinar/atualizar o modelo
    command: ["python", "train.py"]

  api:
    build:
      context: .
      dockerfile: Dockerfile.api
    volumes:
      - ./model:/app/model:ro
    ports:
      - "8000:8000"
    # você sobe esse serviço depois de já ter gerado o model.onnx

6. Como rodar tudo na prática

Dentro da pasta ml-onnx-example/:

1️⃣ Treinar e gerar o modelo ONNX

docker compose run --rm trainer

Isso vai:
- construir a imagem de treino,
- rodar o train.py,
- gravar ./model/model.onnx no host (via volume).

2️⃣ Subir a API

docker compose up api

A API vai:
- carregar model/model.onnx,
- expor http://localhost:8000.

3️⃣ Testar a predição

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"features": [5.1, 3.5, 1.4, 0.2]}'

Resposta esperada (algo nessa linha):

{
  "class_index": 0,
  "class_name": "setosa",
  "probabilities": [0.98, 0.01, 0.01]
}

7. Próximos passos (pensando no seu stack)

Depois que isso estiver redondo, você pode:

Trocar a API Python por uma API em Go
- Manter o model.onnx exatamente igual
- Usar Go apenas para:
  - receber JSON
  - transformar em []float32
  - chamar ONNX Runtime
  - devolver o resultado
Plug-and-play no seu ecossistema
- Colocar essa API como mais um microserviço na sua arquitetura com:
  - Prometheus
  - Grafana
  - Loki
  - Tempo
  - ArgoCD
  - etc.

Ver depois:
03 - Caso prático usando Golang

Exemplo prático

1. Estrutura de pastas do projeto

2. Script de treino em Python (gera model.onnx)

ml/requirements.txt

ml/train.py