Proposta técnica para o time de Engenharia e Data Science da Tractian. Resolve o problema de Cold Start em novos ativos industriais usando Vision Transformers (ViT) + Masked Autoencoders (MAE) + Few-Shot Learning.
Modelos tradicionais de manutenção preditiva exigem semanas de dados históricos para fornecer diagnósticos confiáveis. Clientes recém-onboardados ficam "cegos" a anomalias reais durante esse período — aumentando o risco de churn e reduzindo o valor percebido do produto.
Um Foundation Model pré-treinado em dados históricos de milhares de máquinas que generaliza para novos ativos com menos de 10 minutos de dados, usando Prototypical Networks zero-shot no espaço latente do encoder.
Novo ativo instalado → 5 min de dados → POST /cold-start → Diagnóstico imediato
┌─────────────────────────────────────────────────────────────────┐
│ Docker Compose │
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ pipeline │ │ trainer │ │ api │ │ mlflow │ │
│ │ │ │ │ │ │ │ │ │
│ │ PySpark │───▶│ ViT-MAE │───▶│ FastAPI │ │ :5000 │ │
│ │ STFT │ │ Lightning│ │ ONNX RT │ │ SQLite │ │
│ │ Parquet │ │ FAISS │ │ :8000 │ │ │ │
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │
│ │ │ │ │ │
│ └───────────────┴───────────────┘ │ │
│ ./data/ (volume) ./data/mlflow/ │
└─────────────────────────────────────────────────────────────────┘
| Componente | Tecnologia | Justificativa |
|---|---|---|
| ETL de sinais | PySpark + STFT (torchaudio) | Processamento distribuído; STFT captura harmônicas de falha |
| Foundation Model | ViT-Small/16 (timm) + MAE decoder | Aprendizado auto-supervisionado sem rótulos massivos |
| Pré-treinamento | Masked Autoencoders (75% mask ratio) | Força representações latentes densas da "física" da máquina |
| Inferência cold-start | Prototypical Networks (zero-shot) | Diagnóstico imediato sem retreino |
| Fine-tuning | LinearProbe sobre encoder congelado | Few-shot com 5-10 exemplos por classe |
| Orquestração | Docker Compose | Execução reproduzível em qualquer máquina |
| Rastreio | MLflow | Métricas, artefatos e comparação de experimentos |
| API | FastAPI + ONNX Runtime | Inferência otimizada, Swagger automático |
tractian-cold-start/
├── docker-compose.yml # Orquestra todos os serviços
├── docker-compose.gpu.yml # Override com suporte CUDA
├── .env.example # Variáveis de ambiente
├── Makefile # Atalhos de execução
│
├── services/
│ ├── pipeline/ # ETL: CSV → Espectrogramas
│ │ ├── Dockerfile
│ │ ├── requirements.txt
│ │ └── src/
│ │ ├── main.py # Entrypoint
│ │ ├── signal_processor.py # STFT, normalização, resampling
│ │ ├── spark_job.py # Job PySpark (local mode)
│ │ └── dataset.py # PyTorch Dataset
│ │
│ ├── trainer/ # Pré-treino MAE + Fine-tuning
│ │ ├── Dockerfile
│ │ ├── requirements.txt
│ │ └── src/
│ │ ├── main.py # CLI: --mode pretrain | finetune
│ │ ├── trainer.py # LightningModule
│ │ ├── export_onnx.py # Exporta encoder para ONNX
│ │ └── models/
│ │ ├── vit_mae.py # ViT-Small encoder + MAE decoder
│ │ └── few_shot.py # LinearProbe + PrototypicalClassifier
│ │
│ ├── api/ # Inferência em produção
│ │ ├── Dockerfile
│ │ ├── requirements.txt
│ │ └── src/
│ │ ├── main.py # FastAPI app
│ │ ├── inference.py # ONNX Runtime + Prototypical Networks
│ │ ├── schemas.py # Pydantic models
│ │ └── routers/
│ │ ├── cold_start.py # POST /cold-start
│ │ └── predict.py # POST /predict
│ │
│ └── mlflow/ # Rastreio de experimentos
│ └── Dockerfile
│
├── data/
│ ├── raw/ # CSVs de entrada (montar aqui)
│ ├── spectrograms/ # Espectrogramas gerados (.npy + Parquet)
│ └── models/ # Checkpoints .ckpt + encoder.onnx
│
├── notebooks/
│ └── demo.ipynb # Demo end-to-end com visualizações
│
└── tests/
├── conftest.py
├── test_signal_processor.py
├── test_vit_mae.py
├── test_few_shot.py
└── test_api.py
- Docker >= 24.0
- Docker Compose >= 2.20
- 8 GB de RAM (mínimo para o trainer)
- (Opcional) GPU NVIDIA com nvidia-container-toolkit
git clone <repo-url>
cd tractian-cold-start
# Copia configurações e cria diretórios
make setup
# Build de todas as imagens (pode levar alguns minutos na primeira vez)
make buildColoque seus arquivos CSV em data/raw/. Cada arquivo deve seguir o formato:
timestamp_ms,acceleration_x,acceleration_y,acceleration_z,sample_rate_hz,asset_id,label
0,0.021,-0.003,1.002,12800,motor_001,normal
78,0.023,-0.001,0.998,12800,motor_001,normal
156,0.019,-0.005,1.001,12800,motor_001,bearing_fault| Campo | Tipo | Descrição |
|---|---|---|
timestamp_ms |
int | Timestamp em milissegundos |
acceleration_x/y/z |
float | Aceleração em g por eixo |
sample_rate_hz |
int | Taxa de amostragem original |
asset_id |
string | Identificador único do ativo |
label |
string | Rótulo de falha (opcional — apenas para treinamento supervisionado) |
Rótulos suportados: normal, bearing_fault, imbalance, misalignment, looseness
Os sinais abaixo ilustram as diferenças entre cada condição no domínio do tempo:
make pipeline
# Processa todos os CSVs em data/raw/ → espectrogramas em data/spectrograms/
# Para processar apenas um ativo específico (cold-start rápido):
docker compose run --rm pipeline python src/main.py --asset-id motor_001O pipeline converte os sinais em espectrogramas STFT (224×224px), permitindo que o modelo identifique padrões espectrais distintos por tipo de falha:
make pretrain
# Treina o MAE por 50 épocas (configurável em .env)
# Checkpoints salvos em data/models/
# Métricas disponíveis em http://localhost:5000 (MLflow)O encoder aprende representações que separam as classes no espaço latente (projeção t-SNE):
make finetune
# Treina LinearProbe + calcula centróides Prototypical Networks
# Salva data/models/prototypes.npy e data/models/label_map.jsonmake export-onnx
# Salva data/models/encoder.onnxmake api
# API disponível em http://localhost:8000
# Swagger UI em http://localhost:8000/docs
# MLflow em http://localhost:5000Verifica status dos modelos carregados.
{
"status": "ok",
"onnx_model_loaded": true,
"faiss_index_loaded": true,
"faiss_index_size": 1200,
"embed_dim": 384,
"known_labels": ["normal", "bearing_fault", "imbalance", "misalignment"]
}Diagnóstico imediato para novos ativos (zero-shot, sem retreino).
curl -X POST http://localhost:8000/cold-start \
-H "Content-Type: application/json" \
-d '{
"asset_id": "new_motor_042",
"csv_data": "timestamp_ms,acceleration_x,acceleration_y,acceleration_z,sample_rate_hz\n0,0.031,-0.012,0.992,12800\n..."
}'Resposta:
{
"asset_id": "new_motor_042",
"mode": "prototypical_zero_shot",
"segments_analyzed": 5,
"results": [
{
"segment_id": 0,
"predicted_label": "bearing_fault",
"confidence": 0.87,
"is_anomaly": true
}
],
"dominant_diagnosis": "bearing_fault",
"overall_confidence": 0.84,
"inference_latency_ms": 18.3,
"cold_start_warning": "Diagnóstico baseado em k-NN zero-shot. Precisão aumenta com mais dados."
}Diagnóstico para ativos catalogados (com mais histórico).
Mesmo formato de request/response que /cold-start, sem o campo cold_start_warning.
# Pré-treino com GPU (requer nvidia-container-toolkit)
make pretrain-gpu
# Fine-tuning com GPU
make finetune-gpumake test
# Executa pytest dentro do container trainer
# Cobertura: signal_processor, vit_mae, few_shot, APIResultado esperado:
tests/test_signal_processor.py ........ (8 testes)
tests/test_vit_mae.py ........... (11 testes)
tests/test_few_shot.py .......... (10 testes)
tests/test_api.py ............. (13 testes)
# Instala dependências localmente (ou use o container)
pip install jupyter matplotlib scikit-learn
jupyter notebook notebooks/demo.ipynbO notebook demonstra:
- Geração de sinais sintéticos com falhas conhecidas (BPFO, desbalanceamento, etc.)
- Visualização dos espectrogramas STFT por tipo de falha
- Extração de embeddings via ViT-Small
- Visualização t-SNE do espaço latente
- Curva de acurácia × número de exemplos (KPI principal do Cold Start)
- Chamada end-to-end à API
Edite .env para ajustar os hiperparâmetros:
# Modelo
VIT_MODEL=vit_small_patch16_224
MASK_RATIO=0.75 # fração de patches mascarados no MAE
EMBED_DIM=384 # dimensão dos embeddings ViT-Small
# Treinamento
PRETRAIN_EPOCHS=50
BATCH_SIZE=32
LEARNING_RATE=1.5e-4
PRECISION=16-mixed # FP16 mixed precision (requer GPU)
# Few-Shot / Prototypical Networks
PROTOTYPES_PATH=/app/data/models/prototypes.npy
# Sinal
SAMPLE_RATE=12800 # Hz — taxa de reamostragem
N_FFT=512 # pontos da FFT
HOP_LENGTH=128 # deslocamento entre janelas STFT| Métrica | Meta (PRD) | Como medir |
|---|---|---|
| Time-to-Value | < 1 dia | Tempo entre instalação e primeiro diagnóstico |
| Cold Start Accuracy | > 80% com < 10 min de dados | notebooks/demo.ipynb — curva accuracy × shots |
| Reconstruction Loss | Convergência estável | MLflow → experimento tractian-mae-pretrain |
| FewShotAccuracy@5 | Monitorado por época | MLflow → experimento tractian-few-shot-finetune |
| Inference Latency | < threshold definido | Campo inference_latency_ms na resposta da API |
A Prototypical Networks atinge 90,6% de acurácia com apenas 1 shot, superando o KPI de 80%:
- Masked Autoencoders Are Scalable Vision Learners — He et al., 2022
- An Image is Worth 16x16 Words — Dosovitskiy et al., 2020
- timm — PyTorch Image Models
- FAISS — Facebook AI Similarity Search
- PyTorch Lightning