Como usar o Weights & Biases para rastreamento de experimentos de ML

1. Introdução ao Weights & Biases e sua importância no ciclo de ML

O rastreamento manual de experimentos de Machine Learning é um dos maiores gargalos para equipes que buscam reprodutibilidade e comparação objetiva entre abordagens. Sem uma ferramenta adequada, métricas ficam espalhadas em planilhas, hiperparâmetros são perdidos, e a capacidade de reconstruir um resultado se torna quase impossível após algumas semanas.

O Weights & Biases (W&B) surge como uma plataforma centralizada que resolve esse problema. Suas funcionalidades principais incluem:

  • Rastreamento de experimentos: logging automático de hiperparâmetros, métricas e saídas do modelo
  • Artefatos: versionamento de datasets, checkpoints e modelos serializados
  • Model Registry: gerenciamento do ciclo de vida de modelos em produção
  • Relatórios interativos: criação de dashboards compartilháveis com a equipe

Comparado a alternativas como MLflow e TensorBoard, o W&B se destaca pela facilidade de uso, integração nativa com frameworks modernos e visualizações prontas para uso. Enquanto o TensorBoard exige configuração manual de sumarizadores e o MLflow tem curva de aprendizado maior para recursos avançados, o W&B oferece uma API intuitiva que funciona em menos de 5 minutos.

2. Configuração inicial e integração com projetos de ML

A instalação é simples. Primeiro, instale o pacote e crie sua conta gratuita em wandb.ai:

pip install wandb

Em seguida, faça o login usando sua chave de API (disponível no site após o cadastro):

wandb login

Para iniciar um experimento, use o bloco wandb.init() no início do seu script:

import wandb

wandb.init(
    project="classificacao-imagens",
    name="experimento-resnet50-v1",
    tags=["resnet", "imagenet-pretrained", "lote-1"],
    notes="Primeiro teste com augmentação de dados"
)

Boas práticas de organização: Crie projetos separados por área de problema (ex: classificacao-imagens, nlp-sentimentos). Dentro de cada projeto, use nomes descritivos para runs e adicione tags para filtragem. Para experimentos hierárquicos, utilize o parâmetro group para agrupar runs relacionados.

3. Rastreamento de hiperparâmetros e métricas

O logging de configurações é feito através do objeto wandb.config:

config = wandb.config
config.learning_rate = 0.001
config.batch_size = 64
config.epochs = 50
config.optimizer = "adam"
config.data_path = "./data/train"

Para métricas em tempo real, utilize wandb.log() dentro do loop de treinamento. A função aceita dicionários, permitindo logar múltiplas métricas simultaneamente:

for epoch in range(config.epochs):
    train_loss, train_acc = treinar_modelo()
    val_loss, val_acc = validar_modelo()

    wandb.log({
        "epoch": epoch,
        "train/loss": train_loss,
        "train/accuracy": train_acc,
        "val/loss": val_loss,
        "val/accuracy": val_acc,
        "learning_rate": scheduler.get_last_lr()[0]
    })

Para logs em nível de batch (útil para monitoramento granular), mantenha um contador global:

global_step = 0
for batch in dataloader:
    loss = treinar_batch(batch)
    if global_step % 100 == 0:
        wandb.log({"batch/loss": loss, "global_step": global_step})
    global_step += 1

4. Visualização de gráficos e artefatos avançados

O W&B gera automaticamente gráficos de linha para todas as métricas logadas. Para criar visualizações customizadas, use o painel interativo no site. Você pode combinar runs, adicionar médias móveis e criar superfícies de hiperparâmetros.

Para logging de artefatos, como datasets processados ou checkpoints:

# Salvando um checkpoint
artifact = wandb.Artifact(
    name="modelo-epoch-50",
    type="model",
    description="Checkpoint após 50 épocas de treinamento"
)
artifact.add_file("checkpoint.pth")
wandb.log_artifact(artifact)

# Salvando dataset processado
data_artifact = wandb.Artifact("dataset-aumentado", type="dataset")
data_artifact.add_dir("./data/augmented")
wandb.log_artifact(data_artifact)

Para registrar imagens, tabelas e histogramas:

# Imagens de predição
wandb.log({"predicoes": [wandb.Image(img, caption=f"Classe: {pred}") 
                         for img, pred in zip(imagens_amostra, predicoes)]})

# Tabela de resultados
table = wandb.Table(columns=["id", "real", "predito", "confianca"])
for i, (r, p, c) in enumerate(zip(rotulos_reais, predicoes, confiancas)):
    table.add_data(i, r, p, c)
wandb.log({"tabela_resultados": table})

# Histograma de pesos
for name, param in modelo.named_parameters():
    wandb.log({f"weights/{name}": wandb.Histogram(param.detach().numpy())})

5. Rastreamento de modelos e versionamento de experimentos

O W&B Model Registry permite versionar e gerenciar modelos em diferentes estágios (staging, production):

# Salvando modelo completo
torch.save(modelo.state_dict(), "modelo_final.pth")
artifact = wandb.Artifact("modelo-classificador", type="model")
artifact.add_file("modelo_final.pth")
artifact.add_metadata({"accuracy": 0.94, "framework": "pytorch"})
wandb.log_artifact(artifact)

Para comparar runs lado a lado, acesse a interface web e use a tabela de experimentos. Filtre por tags, métricas mínimas/máximas e visualize gráficos sobrepostos.

Para recuperar checkpoints anteriores:

# Recuperando artefato específico
artifact = wandb.use_artifact("projeto/modelo-classificador:v5")
artifact_dir = artifact.download()
checkpoint_path = f"{artifact_dir}/modelo_final.pth"
modelo.load_state_dict(torch.load(checkpoint_path))

6. Automação e integração com pipelines de ML

Para scripts distribuídos com DDP, cada processo deve inicializar seu próprio run:

import torch.distributed as dist

if dist.get_rank() == 0:  # Apenas rank 0 loga
    wandb.init(project="treinamento-distribuido")
else:
    wandb.init(mode="disabled")  # Desabilita logging nos outros ranks

Integração com frameworks populares:

PyTorch Lightning:

from pytorch_lightning.loggers import WandbLogger

wandb_logger = WandbLogger(project="lightning-project")
trainer = Trainer(logger=wandb_logger, max_epochs=50)

Hugging Face Transformers:

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    report_to="wandb",
    run_name="hf-experimento-1",
    logging_steps=10,
)

Para busca automatizada de hiperparâmetros com wandb.sweep:

sweep_config = {
    "method": "bayes",
    "metric": {"name": "val/accuracy", "goal": "maximize"},
    "parameters": {
        "learning_rate": {"min": 1e-5, "max": 1e-2},
        "batch_size": {"values": [16, 32, 64]},
        "dropout": {"min": 0.1, "max": 0.5}
    }
}

sweep_id = wandb.sweep(sweep_config, project="hyperparameter-search")
wandb.agent(sweep_id, function=treinar_com_sweep, count=20)

7. Boas práticas e troubleshooting

Evitando logs excessivos: Para métricas logadas a cada batch, use wandb.log() a cada N passos. Para artefatos grandes, considere comprimir ou amostrar dados antes do upload.

Gerenciamento de chaves de API: Nunca hardcode a API key no código. Use variáveis de ambiente:

export WANDB_API_KEY=sua_chave_aqui

Em ambientes de produção, utilize secrets managers ou arquivos .env.

Debugging comum:
- Runs órfãos: Sempre feche o run com wandb.finish() ao final do script
- Conflitos de nome: Use nomes únicos para runs ou ative nomes automáticos com wandb.init(name=wandb.util.generate_id())
- Erros de conexão: Configure wandb.init(settings=wandb.Settings(console="off")) para ambientes sem internet, ou use mode="offline" e sincronize depois com wandb sync

Para ambientes offline, colete logs localmente:

wandb.init(mode="offline", project="experimento-offline")
# ... treinamento ...
wandb.finish()
# Depois, sincronize: wandb sync wandb/offline-run-XXX

O Weights & Biases transforma o rastreamento de experimentos de uma tarefa manual e propensa a erros em um processo automatizado e colaborativo. Com as práticas descritas neste artigo, você pode garantir reprodutibilidade, comparar abordagens de forma objetiva e acelerar o ciclo de desenvolvimento de modelos de Machine Learning.

Referências