03 - Caso prático usando Golang
Continuação do artigo:
02 - Caso prático usando Python para treinar um modelo simples (Iris) e exportar para ONNX
Vamos trocar aquela API Python por uma API em Go + chi servindo o mesmo model.onnx.
Vou assumir que você já está gerando o model/model.onnx com o container de treino em Python como montamos antes.
1. Estrutura de pastas (com Go na API)
Mantendo o projeto anterior:
ml-onnx-example/
ml/
train.py
requirements.txt
api-go/
go.mod
main.go
model/
model.onnx # gerado pelo trainer em Python
Dockerfile.ml # já feito antes (treino Python)
Dockerfile.api-go # NOVO: API em Go
docker-compose.yml # vamos ajustar para usar a API Go
2. Módulo Go (api-go/go.mod)
module github.com/seuuser/ml-onnx-example/api-go
go 1.22
require (
github.com/go-chi/chi/v5 v5.0.12
github.com/yalue/onnxruntime_go v0.3.4 // versão de exemplo, ajuste se precisar
)
👀 Obs.:
O wrapper que vou usar é o
github.com/yalue/onnxruntime_go.Se na sua máquina o módulo estiver com outro path (ex.:
github.com/yalue/onnxruntime), é só ajustar o import.
3. Código da API em Go (api-go/main.go)
API simples em chi, com:
-
GET /healthz -
POST /predictrecebendofeatures: [4]float32 -
chamando o
model.onnxvia ONNX Runtime
package main
import (
"encoding/json"
"log"
"math"
"net/http"
"os"
"time"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
ort "github.com/yalue/onnxruntime_go"
)
var (
classNames = []string{"setosa", "versicolor", "virginica"}
modelPath = "model/model.onnx" // dentro do container
)
// structs de request/response
type PredictRequest struct {
Features []float32 `json:"features"`
}
type PredictResponse struct {
ClassIndex int `json:"class_index"`
ClassName string `json:"class_name"`
Probabilities []float32 `json:"probabilities"`
}
// softmax para logits -> probabilidades
func softmax(logits []float32) []float32 {
if len(logits) == 0 {
return []float32{}
}
maxLogit := logits[0]
for _, v := range logits[1:] {
if v > maxLogit {
maxLogit = v
}
}
var sum float64
exps := make([]float64, len(logits))
for i, v := range logits {
ev := math.Exp(float64(v - maxLogit))
exps[i] = ev
sum += ev
}
out := make([]float32, len(logits))
if sum == 0 {
return out
}
for i, ev := range exps {
out[i] = float32(ev / sum)
}
return out
}
// handler /healthz
func healthzHandler(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
// handler /predict
func predictHandler(w http.ResponseWriter, r *http.Request) {
var req PredictRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
writeError(w, http.StatusBadRequest, "JSON inválido")
return
}
if len(req.Features) != 4 {
writeError(w, http.StatusBadRequest, "é obrigatório enviar exatamente 4 features")
return
}
// cria tensor de entrada
inputData := make([]float32, 4)
copy(inputData, req.Features)
inputShape := ort.NewShape(1, 4) // batch=1, 4 features
inputTensor, err := ort.NewTensor(inputShape, inputData)
if err != nil {
log.Printf("erro criando tensor de entrada: %v", err)
writeError(w, http.StatusInternalServerError, "erro interno (tensor input)")
return
}
defer inputTensor.Destroy()
// tensor de saída (1,3)
outputShape := ort.NewShape(1, 3)
outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
if err != nil {
log.Printf("erro criando tensor de saída: %v", err)
writeError(w, http.StatusInternalServerError, "erro interno (tensor output)")
return
}
defer outputTensor.Destroy()
// cria sessão ONNX
session, err := ort.NewSession[float32](
modelPath,
[]string{"input"}, // nome da entrada no export do PyTorch
[]string{"output"}, // nome da saída no export do PyTorch
[]*ort.Tensor[float32]{inputTensor},
[]*ort.Tensor[float32]{outputTensor},
)
if err != nil {
log.Printf("erro criando sessão onnx: %v", err)
writeError(w, http.StatusInternalServerError, "erro interno (sessão onnx)")
return
}
defer session.Destroy()
// executa rede
if err := session.Run(); err != nil {
log.Printf("erro executando sessão: %v", err)
writeError(w, http.StatusInternalServerError, "erro interno (run)")
return
}
// pega logits de saída
outData := outputTensor.GetData()
if len(outData) != 3 {
log.Printf("saida inesperada: %v", outData)
writeError(w, http.StatusInternalServerError, "erro interno (saida)")
return
}
probs := softmax(outData)
// encontra índice da classe com maior prob
maxIdx := 0
maxVal := probs[0]
for i := 1; i < len(probs); i++ {
if probs[i] > maxVal {
maxVal = probs[i]
maxIdx = i
}
}
className := "desconhecida"
if maxIdx >= 0 && maxIdx < len(classNames) {
className = classNames[maxIdx]
}
resp := PredictResponse{
ClassIndex: maxIdx,
ClassName: className,
Probabilities: probs,
}
writeJSON(w, http.StatusOK, resp)
}
// helpers de resposta JSON
func writeJSON(w http.ResponseWriter, status int, payload any) {
w.Header().Set("Content-Type", "application/json; charset=utf-8")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(payload)
}
func writeError(w http.ResponseWriter, status int, msg string) {
writeJSON(w, status, map[string]string{
"error": msg,
})
}
func main() {
// Path da lib onnxruntime (opcional, mas recomendado)
// Exemplo: /usr/local/lib/libonnxruntime.so
if libPath := os.Getenv("ONNXRUNTIME_LIB"); libPath != "" {
ort.SetSharedLibraryPath(libPath)
}
if err := ort.InitializeEnvironment(); err != nil {
log.Fatalf("erro inicializando onnxruntime: %v", err)
}
defer func() {
if err := ort.DestroyEnvironment(); err != nil {
log.Printf("erro destruindo ambiente onnxruntime: %v", err)
}
}()
r := chi.NewRouter()
// middlewares padrão
r.Use(middleware.RequestID)
r.Use(middleware.RealIP)
r.Use(middleware.Logger)
r.Use(middleware.Recoverer)
r.Use(middleware.Timeout(10 * time.Second))
// rotas
r.Get("/healthz", healthzHandler)
r.Post("/predict", predictHandler)
addr := ":8000"
log.Printf("API Go ONNX escutando em %s", addr)
if err := http.ListenAndServe(addr, r); err != nil {
log.Fatalf("erro servidor HTTP: %v", err)
}
}
✅ Pra produção real, o ideal é:
reutilizar sessão e tensores com mutex em vez de criar a cada request,
ou ter um pool de sessões.
Mas pra ficar didático, deixei 1 sessão / request.
4. Dockerfile da API em Go (Dockerfile.api-go)
Aqui a parte mais chata é ter a libonnxruntime.so disponível.
Vou deixar um passo genérico com download da lib; você só ajusta a versão/URL.
FROM golang:1.22-bookworm AS builder
WORKDIR /app
# Copia o código
COPY api-go/go.mod api-go/go.sum ./api-go/
WORKDIR /app/api-go
RUN go mod download
COPY api-go/*.go ./
# Compila binário estático-ish
RUN CGO_ENABLED=1 GOOS=linux GOARCH=amd64 go build -o /app/bin/api-go
# -------------------------------------------------------------------
FROM debian:bookworm-slim
WORKDIR /app
# Dependências mínimas
RUN apt-get update && apt-get install -y \
ca-certificates \
curl \
&& rm -rf /var/lib/apt/lists/*
# === ONNXRUNTIME ===
# ATENÇÃO: ajuste a versão e URL conforme a versão que quiser usar
# Exemplo de URL (ver releases no GitHub microsoft/onnxruntime):
# https://github.com/microsoft/onnxruntime/releases/download/v1.20.0/onnxruntime-linux-x64-1.20.0.tgz
ENV ONNX_VERSION=1.20.0
ENV ONNXRUNTIME_LIB=/usr/local/lib/libonnxruntime.so
RUN curl -L "https://github.com/microsoft/onnxruntime/releases/download/v${ONNX_VERSION}/onnxruntime-linux-x64-${ONNX_VERSION}.tgz" \
-o /tmp/onnxruntime.tgz && \
mkdir -p /opt/onnxruntime && \
tar -xzf /tmp/onnxruntime.tgz -C /opt/onnxruntime --strip-components=1 && \
cp /opt/onnxruntime/lib/libonnxruntime.so /usr/local/lib/ && \
rm /tmp/onnxruntime.tgz
# Copia o binário
COPY --from=builder /app/bin/api-go /usr/local/bin/api-go
# Pasta do modelo (montada como volume pelo docker-compose)
RUN mkdir -p /app/model
WORKDIR /app
EXPOSE 8000
CMD ["api-go"]
5. Ajustando o docker-compose.yml para usar a API em Go
Reaproveitando o serviço trainer (Python) que já tínhamos, vamos trocar a api Python por api-go.
version: "3.9"
services:
trainer:
build:
context: .
dockerfile: Dockerfile.ml
volumes:
- ./model:/app/model
command: ["python", "train.py"]
api-go:
build:
context: .
dockerfile: Dockerfile.api-go
volumes:
- ./model:/app/model:ro
environment:
- ONNXRUNTIME_LIB=/usr/local/lib/libonnxruntime.so
ports:
- "8000:8000"
6. Fluxo para rodar tudo
Dentro de ml-onnx-example/:
1️⃣ Treinar o modelo (Python → ONNX)
docker compose run --rm trainer
Gera ./model/model.onnx.
2️⃣ Subir API Go
docker compose up api-go --build
3️⃣ Chamar 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 (mais ou menos):
{
"class_index": 0,
"class_name": "setosa",
"probabilities": [0.98, 0.01, 0.01]
}
Ver depois:
04 - Convertendo para arquitetura hexagonal