318 lines
13 KiB
Python
318 lines
13 KiB
Python
import json
|
|
import os
|
|
import time
|
|
import argparse
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader
|
|
from fast_scnn import FastSCNN
|
|
from roi_seg_dataset import ROISegDataset
|
|
import matplotlib.pyplot as plt
|
|
|
|
# ⚙️ Configurações
|
|
with open("config.json", "r") as f:
|
|
config = json.load(f)
|
|
MODELO = config["camera"]
|
|
MODEL_NAME = config["model_name"]
|
|
RESOLUCAO = config["resolucao"]
|
|
ROI_INICIO = config["roi_inicio"]
|
|
ROI_TAMANHO = config["roi_tamanho"]
|
|
MAIN_CLASS_NAME = config["main_class_name"]
|
|
save_path = os.path.join(MODELO, "backup", config["modelo"], MODEL_NAME)
|
|
dataset_path = os.path.join(MODELO, "dataset")
|
|
labelmap_path = os.path.join(dataset_path, "labelmap.txt")
|
|
batch_size = 8
|
|
num_workers = 4
|
|
|
|
# ---- Helpers de métricas ----
|
|
@torch.no_grad()
|
|
def confmat_update(confmat, pred, target, num_classes, ignore_index=None):
|
|
# pred, target: (B,H,W)
|
|
if ignore_index is not None:
|
|
mask = target != ignore_index
|
|
target = target[mask]
|
|
pred = pred[mask]
|
|
k = (target * num_classes + pred).to(torch.int64)
|
|
binc = torch.bincount(k, minlength=num_classes**2)
|
|
confmat += binc.reshape(num_classes, num_classes)
|
|
return confmat
|
|
|
|
def metrics_from_confmat(confmat, main_class_id=None):
|
|
# confmat: CxC
|
|
cm = confmat.float()
|
|
tp = torch.diag(cm)
|
|
fp = cm.sum(0) - tp
|
|
fn = cm.sum(1) - tp
|
|
denom_iou = tp + fp + fn + 1e-7
|
|
iou_per_class = tp / denom_iou
|
|
miou = iou_per_class.mean().item()
|
|
pix_acc = tp.sum() / (cm.sum() + 1e-7)
|
|
|
|
main_class_metrics = None
|
|
if main_class_id is not None and 0 <= main_class_id < cm.shape[0]:
|
|
p = tp[main_class_id] / (tp[main_class_id] + fp[main_class_id] + 1e-7)
|
|
r = tp[main_class_id] / (tp[main_class_id] + fn[main_class_id] + 1e-7)
|
|
f1 = 2 * p * r / (p + r + 1e-7)
|
|
main_class_metrics = {
|
|
"precision": p.item(),
|
|
"recall": r.item(),
|
|
"f1": f1.item(),
|
|
"iou": iou_per_class[main_class_id].item(),
|
|
}
|
|
return {
|
|
"miou": miou,
|
|
"pixel_acc": pix_acc.item(),
|
|
"iou_per_class": iou_per_class.cpu().tolist(),
|
|
"main_class": main_class_metrics
|
|
}
|
|
|
|
def train(args):
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Device: {device}")
|
|
|
|
# --- Dataset ---
|
|
ds_train = ROISegDataset(
|
|
os.path.join(dataset_path, "split", "train"),
|
|
save_path, ROI_INICIO, ROI_TAMANHO,
|
|
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
|
)
|
|
ds_val = ROISegDataset(
|
|
os.path.join(dataset_path, "split", "val"),
|
|
save_path, ROI_INICIO, ROI_TAMANHO,
|
|
RESOLUCAO[0], RESOLUCAO[1], labelmap_path
|
|
)
|
|
|
|
dl_train = DataLoader(ds_train, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)
|
|
dl_val = DataLoader(ds_val, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)
|
|
|
|
# Detecta automaticamente o ID da classe ERVA
|
|
main_class_id = None
|
|
try:
|
|
if hasattr(ds_train, "classes") and isinstance(ds_train.classes, dict):
|
|
for k, v in ds_train.classes.items():
|
|
if isinstance(v, str) and MAIN_CLASS_NAME in v.lower():
|
|
main_class_id = k
|
|
break
|
|
elif isinstance(ds_train.classes, (list, tuple)):
|
|
main_class_id = next((i for i, c in enumerate(ds_train.classes) if isinstance(c, str) and MAIN_CLASS_NAME in c.lower()), None)
|
|
|
|
if main_class_id is not None:
|
|
print(f"🌿 Classe PRIMARIA detectada: id={main_class_id}, nome='{ds_train.classes[main_class_id]}'")
|
|
else:
|
|
print("⚠️ Classe PRIMARIA não encontrada; métricas específicas da classe primaria serão puladas.")
|
|
except Exception as e:
|
|
print(f"⚠️ Erro ao detectar classe PRIMARIA: {e}")
|
|
|
|
num_classes = len(ds_train.classes)
|
|
|
|
# --- Modelo / Otimizador / Schedulers ---
|
|
model = FastSCNN(num_classes=num_classes).to(device)
|
|
criterion = nn.CrossEntropyLoss(ignore_index=ds_train.ignore_id)
|
|
optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
|
|
|
|
# Scheduler inteligente: começa em Cosine, muda pra Plateau se travar
|
|
min_lr = getattr(args, "min_lr", 1e-6)
|
|
plateau_factor = getattr(args, "plateau_factor", 0.5)
|
|
plateau_patience = getattr(args, "plateau_patience", 6) # épocas sem melhora antes de trocar
|
|
plateau_cooldown = getattr(args, "plateau_cooldown", 1)
|
|
|
|
cosine = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=min_lr)
|
|
plateau = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
optimizer, mode="min", factor=plateau_factor,
|
|
patience=plateau_patience, cooldown=plateau_cooldown,
|
|
min_lr=min_lr, verbose=True
|
|
)
|
|
active_sched = "cosine"
|
|
|
|
scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
|
|
|
|
start_epoch = 1
|
|
best_val_loss = float("inf")
|
|
best_main_class_f1 = -1.0
|
|
train_loss_history, val_loss_history, lr_history = [], [], []
|
|
f1_history, miou_history = [], []
|
|
|
|
# --- no topo (config) ---
|
|
patience_loss = 12 # ligeiramente > plateau_patience + 2
|
|
patience_f1 = 6 # deixa o F1 respirar
|
|
delta_f1_min = 0.0015 # ignora ruído
|
|
grace_after_switch = 4 # épocas de graça após mudar pro Plateau
|
|
|
|
no_imp_loss = 0
|
|
no_imp_f1 = 0
|
|
epochs_since_switch = 0
|
|
active_sched = "cosine" # como já está
|
|
|
|
# --- Checkpoint ---
|
|
if args.checkpoint and os.path.exists(args.checkpoint):
|
|
print(f"🔁 Carregando modelo salvo: {args.checkpoint}")
|
|
checkpoint = torch.load(args.checkpoint, map_location=device)
|
|
if "model" in checkpoint:
|
|
model.load_state_dict(checkpoint["model"])
|
|
optimizer.load_state_dict(checkpoint["optimizer"])
|
|
scaler.load_state_dict(checkpoint["scaler"])
|
|
start_epoch = checkpoint.get("epoch", 1) + 1
|
|
best_val_loss = checkpoint.get("best_val_loss", float("inf"))
|
|
else:
|
|
model.load_state_dict(checkpoint)
|
|
|
|
# --- Loop de treino ---
|
|
for epoch in range(start_epoch, args.epochs + 1):
|
|
t0 = time.time()
|
|
|
|
# ----- Treino -----
|
|
model.train()
|
|
running_train_loss = 0
|
|
for x, y in dl_train:
|
|
x, y = x.to(device), y.to(device)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
with torch.cuda.amp.autocast(enabled=args.amp):
|
|
logits = model(x)
|
|
loss = criterion(logits, y)
|
|
scaler.scale(loss).backward()
|
|
scaler.step(optimizer)
|
|
scaler.update()
|
|
running_train_loss += loss.item() * x.size(0)
|
|
|
|
avg_train_loss = running_train_loss / len(ds_train)
|
|
train_loss_history.append(avg_train_loss)
|
|
|
|
# ----- Validação + métricas -----
|
|
model.eval()
|
|
running_val_loss = 0
|
|
confmat = torch.zeros((num_classes, num_classes), dtype=torch.int64, device=device)
|
|
|
|
with torch.no_grad():
|
|
for x, y in dl_val:
|
|
x, y = x.to(device), y.to(device)
|
|
with torch.cuda.amp.autocast(enabled=args.amp):
|
|
logits = model(x)
|
|
loss = criterion(logits, y)
|
|
running_val_loss += loss.item() * x.size(0)
|
|
|
|
pred = logits.argmax(1)
|
|
confmat = confmat_update(confmat, pred, y, num_classes, ignore_index=ds_train.ignore_id)
|
|
|
|
avg_val_loss = running_val_loss / len(ds_val)
|
|
val_loss_history.append(avg_val_loss)
|
|
|
|
m = metrics_from_confmat(confmat, main_class_id=main_class_id)
|
|
miou_history.append(m["miou"])
|
|
main_class_f1 = m["main_class"]["f1"] if (m["main_class"] is not None) else None
|
|
if main_class_f1 is not None:
|
|
f1_history.append(main_class_f1)
|
|
cur_lr = optimizer.param_groups[0]["lr"]
|
|
lr_history.append(cur_lr)
|
|
|
|
elapsed = time.time() - t0
|
|
msg = (f"[{epoch}/{args.epochs}] "
|
|
f"train_loss={avg_train_loss:.4f} "
|
|
f"val_loss={avg_val_loss:.4f} "
|
|
f"mIoU={m['miou']:.4f} "
|
|
f"pixAcc={m['pixel_acc']:.4f} "
|
|
f"lr={cur_lr:.2e} "
|
|
f"time={elapsed:.1f}s")
|
|
if main_class_f1 is not None:
|
|
msg += f" | {MAIN_CLASS_NAME}: F1={main_class_f1:.4f} IoU={m['main_class']['iou']:.4f}"
|
|
print(msg)
|
|
|
|
# ----- Tracking de melhora por LOSS -----
|
|
improved_loss = avg_val_loss < best_val_loss - 1e-6
|
|
if improved_loss:
|
|
best_val_loss = avg_val_loss
|
|
no_imp_loss = 0
|
|
# checkpoint por loss
|
|
torch.save(model.state_dict(), os.path.join(save_path, f"{MODEL_NAME}_best.pth"))
|
|
torch.save({
|
|
"model": model.state_dict(),
|
|
"optimizer": optimizer.state_dict(),
|
|
"scaler": scaler.state_dict(),
|
|
"epoch": epoch,
|
|
"best_val_loss": best_val_loss
|
|
}, os.path.join(save_path, f"{MODEL_NAME}_best_checkpoint.pth"))
|
|
print("✅ Novo melhor modelo salvo (val_loss).")
|
|
else:
|
|
no_imp_loss += 1
|
|
|
|
# ----- Tracking + checkpoint por F1 da classe principal -----
|
|
if main_class_f1 is not None:
|
|
if main_class_f1 > best_main_class_f1 + delta_f1_min:
|
|
best_main_class_f1 = main_class_f1
|
|
no_imp_f1 = 0
|
|
torch.save(model.state_dict(), os.path.join(save_path, f"{MODEL_NAME}_best_f1_{MAIN_CLASS_NAME}.pth"))
|
|
print(f"🌿💾 Checkpoint salvo (melhor F1 da {MAIN_CLASS_NAME}).")
|
|
else:
|
|
no_imp_f1 += 1
|
|
else:
|
|
# se não houver F1 (ex: id não definido), ignora o critério
|
|
no_imp_f1 = 0
|
|
|
|
# ----- Scheduler inteligente -----
|
|
if active_sched == "cosine":
|
|
# se travar por plateau_patience, troca pra ReduceLROnPlateau
|
|
if no_imp_loss >= plateau_patience:
|
|
active_sched = "plateau"
|
|
print("🔁 Mudando scheduler: Cosine → ReduceLROnPlateau (platô detectado).")
|
|
# resets ao trocar
|
|
no_imp_loss = 0
|
|
no_imp_f1 = 0
|
|
epochs_since_switch = 0
|
|
plateau.step(avg_val_loss) # primeiro passo do plateau
|
|
# (opcional) “adiantar” a queda do LR:
|
|
for g in optimizer.param_groups:
|
|
g['lr'] = max(g['lr'] * plateau_factor, min_lr)
|
|
else:
|
|
cosine.step()
|
|
else:
|
|
plateau.step(avg_val_loss)
|
|
epochs_since_switch += 1
|
|
|
|
# ----- Log de estagnação -----
|
|
print(f"⏳ Sem melhora — loss: {no_imp_loss}/{patience_loss}, {MAIN_CLASS_NAME}: {no_imp_f1}/{patience_f1}")
|
|
|
|
# ----- Early stopping bi-critério (com 'graça' após switch) -----
|
|
if (no_imp_loss >= patience_loss and
|
|
(main_class_f1 is None or no_imp_f1 >= patience_f1) and
|
|
(active_sched == "cosine" or epochs_since_switch >= grace_after_switch)):
|
|
print("⏹ Early stopping: loss e F1 sem melhora (com período de graça respeitado).")
|
|
break
|
|
|
|
# ----- Plots periódicos -----
|
|
if epoch % 5 == 0 or epoch == args.epochs:
|
|
x_epochs = list(range(start_epoch, start_epoch + len(train_loss_history)))
|
|
# Loss
|
|
plt.figure()
|
|
plt.plot(x_epochs, train_loss_history, marker="o", label="Train Loss")
|
|
plt.plot(x_epochs, val_loss_history, marker="s", label="Val Loss")
|
|
plt.xlabel("Época"); plt.ylabel("Loss"); plt.grid(True); plt.legend(); plt.title("Curva de Loss")
|
|
plt.tight_layout()
|
|
plt.savefig(os.path.join(save_path, "loss_curve.png")); plt.close()
|
|
# LR
|
|
plt.figure()
|
|
plt.plot(x_epochs, lr_history, marker=".")
|
|
plt.xlabel("Época"); plt.ylabel("LR"); plt.grid(True); plt.title("Learning Rate")
|
|
plt.tight_layout()
|
|
plt.savefig(os.path.join(save_path, "lr_curve.png")); plt.close()
|
|
# mIoU e F1(erva)
|
|
plt.figure()
|
|
plt.plot(x_epochs, miou_history, marker="^", label="mIoU")
|
|
if len(f1_history) == len(miou_history):
|
|
plt.plot(x_epochs, f1_history, marker="*", label=f"F1 {MAIN_CLASS_NAME}")
|
|
plt.xlabel("Época"); plt.ylabel("Score"); plt.grid(True); plt.legend(); plt.title(f"mIoU / F1({MAIN_CLASS_NAME})")
|
|
plt.tight_layout()
|
|
plt.savefig(os.path.join(save_path, "metrics_curve.png")); plt.close()
|
|
|
|
def parse_args():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--epochs", type=int, default=30)
|
|
ap.add_argument("--lr", type=float, default=3e-4)
|
|
ap.add_argument("--amp", action="store_true")
|
|
ap.add_argument("--checkpoint", type=str, default=None, help="Caminho do modelo .pth para continuar o treinamento")
|
|
return ap.parse_args()
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
os.makedirs(save_path, exist_ok=True)
|
|
train(args)
|