| """
|
| Implementação concreta (PyTorch) do modelo treinado pelo pipeline.
|
|
|
| Arquitetura: MLP simples para dados tabulares. A configuração de
|
| camadas, dropout, otimizador, etc. vem de ``config`` (lido do
|
| ``config.json`` do treino) ou de defaults razoáveis.
|
|
|
| A classe ``TabularMLP`` herda de ``BaseModel`` (em ``docker/base_model.py``)
|
| para manter a API uniforme entre frameworks. Uma função ``create_model``
|
| fábrica é exposta no fim do arquivo para que ``main.py`` continue
|
| funcionando sem mudanças.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import os
|
| import sys
|
| from typing import Any, Dict, List, Optional
|
|
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
| import torch.optim as optim
|
| from torch.utils.data import DataLoader
|
|
|
|
|
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
| from base_model import BaseModel
|
| from utils import emit_progress
|
|
|
|
|
|
|
|
|
|
|
| class _MLP(nn.Module):
|
| def __init__(
|
| self,
|
| input_size: int,
|
| hidden_dims: List[int],
|
| num_classes: int,
|
| dropout: float = 0.0,
|
| ) -> None:
|
| super().__init__()
|
| layers: List[nn.Module] = []
|
| prev = input_size
|
| for h in hidden_dims:
|
| layers.append(nn.Linear(prev, h))
|
| layers.append(nn.ReLU())
|
| if dropout and dropout > 0:
|
| layers.append(nn.Dropout(dropout))
|
| prev = h
|
| layers.append(nn.Linear(prev, num_classes))
|
| self.net = nn.Sequential(*layers)
|
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| return self.net(x)
|
|
|
|
|
|
|
|
|
|
|
| class TabularMLP(BaseModel):
|
| """MLP tabular em PyTorch, plugável no pipeline."""
|
|
|
| def build_model(self) -> nn.Module:
|
| hp = self.config.get("hyperparameters", {})
|
| hidden_dims = self.config.get("hidden_dims") or hp.get("hidden_dims") or [
|
| self.config.get("hidden_dim1", 32),
|
| self.config.get("hidden_dim2", 16),
|
| ]
|
| dropout = float(hp.get("dropout", self.config.get("dropout", 0.0)))
|
|
|
|
|
| self.device = torch.device(
|
| "cuda"
|
| if torch.cuda.is_available() and self.config.get("use_cuda", True)
|
| else "cpu"
|
| )
|
|
|
| model = _MLP(
|
| input_size=self.input_size,
|
| hidden_dims=list(hidden_dims),
|
| num_classes=self.num_classes,
|
| dropout=dropout,
|
| ).to(self.device)
|
|
|
| self.criterion = nn.CrossEntropyLoss()
|
| self._hidden_dims = list(hidden_dims)
|
| self._dropout = dropout
|
| return model
|
|
|
|
|
| def train_model(
|
| self,
|
| train_loader: DataLoader,
|
| val_loader: Optional[DataLoader],
|
| epochs: int,
|
| lr: float,
|
| **kwargs: Any,
|
| ) -> Dict[str, list]:
|
| hp = self.config.get("hyperparameters", {})
|
| epochs = int(hp.get("epochs", epochs))
|
| lr = float(hp.get("learning_rate", lr))
|
| weight_decay = float(hp.get("weight_decay", 0.0))
|
| optimizer_name = str(hp.get("optimizer", "adam")).lower()
|
|
|
| optimizer = self._build_optimizer(optimizer_name, lr, weight_decay)
|
|
|
| history = {
|
| "train_loss": [], "train_acc": [],
|
| "val_loss": [], "val_acc": [],
|
| }
|
|
|
| n_batches = max(len(train_loader), 1)
|
| total_train_steps = max(epochs * n_batches, 1)
|
| best_val_acc = -float("inf")
|
| best_epoch = -1
|
| for epoch in range(1, epochs + 1):
|
| self.model.train()
|
| run_loss = 0.0
|
| correct = 0
|
| total = 0
|
| for batch_idx, (X, y) in enumerate(train_loader, start=1):
|
| X = X.to(self.device)
|
| y = y.to(self.device)
|
|
|
| optimizer.zero_grad()
|
| logits = self.model(X)
|
| loss = self.criterion(logits, y)
|
| loss.backward()
|
| optimizer.step()
|
|
|
| run_loss += float(loss.item()) * X.size(0)
|
| preds = logits.argmax(dim=1)
|
| correct += int((preds == y).sum().item())
|
| total += int(y.size(0))
|
|
|
| global_step = (epoch - 1) * n_batches
|
| inner_pct = int(global_step * 100 / total_train_steps)
|
| emit_progress(inner_pct, total_train_steps)
|
|
|
| train_loss = run_loss / max(total, 1)
|
| train_acc = correct / max(total, 1)
|
| history["train_loss"].append(round(train_loss, 6))
|
| history["train_acc"].append(round(train_acc, 6))
|
|
|
| val_loss, val_acc = (float("nan"), float("nan"))
|
| if val_loader is not None:
|
| val_metrics = self.evaluate(val_loader)
|
| val_loss = val_metrics["loss"]
|
| val_acc = val_metrics["accuracy"]
|
| if val_acc > best_val_acc:
|
| best_val_acc = val_acc
|
| best_epoch = epoch
|
| history["val_loss"].append(round(val_loss, 6))
|
| history["val_acc"].append(round(val_acc, 6))
|
|
|
| print(
|
| f"Epoch [{epoch:>3}/{epochs}] "
|
| f"train_loss={train_loss:.4f} train_acc={train_acc:.4f} "
|
| f"val_loss={val_loss:.4f} val_acc={val_acc:.4f}"
|
| )
|
|
|
| history["best_epoch"] = best_epoch
|
| history["best_val_acc"] = round(best_val_acc, 6) if best_epoch > 0 else None
|
| self.history = history
|
| return history
|
|
|
|
|
| def evaluate(self, data_loader: DataLoader) -> Dict[str, float]:
|
| self.model.eval()
|
| loss_sum = 0.0
|
| correct = 0
|
| total = 0
|
| with torch.no_grad():
|
| for X, y in data_loader:
|
| X = X.to(self.device)
|
| y = y.to(self.device)
|
| logits = self.model(X)
|
| loss = self.criterion(logits, y)
|
| loss_sum += float(loss.item()) * X.size(0)
|
| preds = logits.argmax(dim=1)
|
| correct += int((preds == y).sum().item())
|
| total += int(y.size(0))
|
| return {
|
| "loss": loss_sum / max(total, 1),
|
| "accuracy": correct / max(total, 1),
|
| "n": total,
|
| }
|
|
|
|
|
| def predict(self, inputs: Any) -> np.ndarray:
|
| self.model.eval()
|
| with torch.no_grad():
|
| if isinstance(inputs, DataLoader):
|
| outs: List[np.ndarray] = []
|
| for batch in inputs:
|
| X = batch[0] if isinstance(batch, (tuple, list)) else batch
|
| X = X.to(self.device)
|
| outs.append(self.model(X).argmax(dim=1).cpu().numpy())
|
| return np.concatenate(outs, axis=0)
|
| if isinstance(inputs, np.ndarray):
|
| X = torch.from_numpy(inputs.astype(np.float32)).to(self.device)
|
| elif isinstance(inputs, torch.Tensor):
|
| X = inputs.to(self.device)
|
| else:
|
| X = torch.tensor(inputs, dtype=torch.float32).to(self.device)
|
| if X.ndim == 1:
|
| X = X.unsqueeze(0)
|
| return self.model(X).argmax(dim=1).cpu().numpy()
|
|
|
|
|
| def save_model(self, filename: str) -> str:
|
| path = os.path.join(self.model_dir, filename)
|
| torch.save(
|
| {
|
| "state_dict": self.model.state_dict(),
|
| "input_size": self.input_size,
|
| "num_classes": self.num_classes,
|
| "hidden_dims": self._hidden_dims,
|
| "dropout": self._dropout,
|
| "project_name": self.project_name,
|
| },
|
| path,
|
| )
|
| print(f"✓ Model saved to {path}")
|
| return path
|
|
|
| def load_model(self, filename: str) -> None:
|
| path = filename if os.path.isabs(filename) else os.path.join(
|
| self.model_dir, filename
|
| )
|
| ckpt = torch.load(path, map_location=self.device)
|
| self._hidden_dims = ckpt.get("hidden_dims", self._hidden_dims)
|
| self._dropout = ckpt.get("dropout", self._dropout)
|
| self.model = _MLP(
|
| input_size=ckpt["input_size"],
|
| hidden_dims=self._hidden_dims,
|
| num_classes=ckpt["num_classes"],
|
| dropout=self._dropout,
|
| ).to(self.device)
|
| self.model.load_state_dict(ckpt["state_dict"])
|
|
|
|
|
| def _build_optimizer(
|
| self, name: str, lr: float, weight_decay: float
|
| ) -> optim.Optimizer:
|
| params = self.model.parameters()
|
| if name == "sgd":
|
| return optim.SGD(params, lr=lr, momentum=0.9, weight_decay=weight_decay)
|
| if name == "rmsprop":
|
| return optim.RMSprop(params, lr=lr, weight_decay=weight_decay)
|
|
|
| return optim.Adam(params, lr=lr, weight_decay=weight_decay)
|
|
|
|
|
|
|
|
|
|
|
| def create_model(
|
| input_size: int,
|
| num_classes: int,
|
| project_name: str,
|
| base_path: str,
|
| config: Optional[Dict[str, Any]] = None,
|
| ) -> TabularMLP:
|
| """Cria um ``TabularMLP`` pronto pra uso pelo main.py."""
|
| return TabularMLP(
|
| input_size=input_size,
|
| num_classes=num_classes,
|
| project_name=project_name,
|
| base_path=base_path,
|
| framework="pytorch",
|
| config=config or {},
|
| )
|
|
|
