1172 lines
38 KiB
Python
1172 lines
38 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
"""
|
||
|
|
_11_validate_onnx.py
|
||
|
|
|
||
|
|
Valida fidelidade entre:
|
||
|
|
- modelo PyTorch .pt
|
||
|
|
- modelo ONNX .onnx
|
||
|
|
|
||
|
|
para o SegFormer OAK-FCC-3 Multi-Head.
|
||
|
|
|
||
|
|
Exemplo:
|
||
|
|
|
||
|
|
python _11_validate_onnx.py --config config.json --max_samples 20 --device cuda --onnx_provider cuda --torch_no_amp
|
||
|
|
|
||
|
|
Para validar o ONNX com saída já redimensionada:
|
||
|
|
|
||
|
|
python _11_validate_onnx.py --config config.json --max_samples 20 --device cuda --onnx_provider cuda --torch_no_amp --compare_at_input_size
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import os
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import copy
|
||
|
|
import argparse
|
||
|
|
import importlib.util
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Dict, List, Optional, Tuple
|
||
|
|
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.nn.functional as F
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"]
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Utils
|
||
|
|
# ============================================================
|
||
|
|
|
||
|
|
def load_json(path: str | Path) -> dict:
|
||
|
|
with open(path, "r", encoding="utf-8") as f:
|
||
|
|
return json.load(f)
|
||
|
|
|
||
|
|
|
||
|
|
def save_json(path: str | Path, data: dict):
|
||
|
|
path = Path(path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with path.open("w", encoding="utf-8") as f:
|
||
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]:
|
||
|
|
if path_like is None:
|
||
|
|
return None
|
||
|
|
|
||
|
|
p = Path(path_like)
|
||
|
|
if p.is_absolute():
|
||
|
|
return p
|
||
|
|
|
||
|
|
if base is None:
|
||
|
|
base = Path.cwd()
|
||
|
|
|
||
|
|
return (base / p).resolve()
|
||
|
|
|
||
|
|
|
||
|
|
def import_train_module(train_script_path: str | Path):
|
||
|
|
train_script_path = Path(train_script_path)
|
||
|
|
|
||
|
|
if not train_script_path.exists():
|
||
|
|
raise FileNotFoundError(f"Script de treino não encontrado: {train_script_path}")
|
||
|
|
|
||
|
|
spec = importlib.util.spec_from_file_location(
|
||
|
|
"train_multihead_module",
|
||
|
|
str(train_script_path.resolve())
|
||
|
|
)
|
||
|
|
|
||
|
|
if spec is None or spec.loader is None:
|
||
|
|
raise RuntimeError(f"Não consegui importar o script: {train_script_path}")
|
||
|
|
|
||
|
|
module = importlib.util.module_from_spec(spec)
|
||
|
|
spec.loader.exec_module(module)
|
||
|
|
return module
|
||
|
|
|
||
|
|
|
||
|
|
def softmax_np(logits: np.ndarray, axis: int = 1) -> np.ndarray:
|
||
|
|
x = logits.astype(np.float32)
|
||
|
|
x = x - np.max(x, axis=axis, keepdims=True)
|
||
|
|
e = np.exp(x)
|
||
|
|
return e / np.clip(np.sum(e, axis=axis, keepdims=True), 1e-12, None)
|
||
|
|
|
||
|
|
|
||
|
|
def resize_logits_np_nchw(logits: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
|
||
|
|
"""
|
||
|
|
logits: [N,C,H,W]
|
||
|
|
target_hw: (H,W)
|
||
|
|
"""
|
||
|
|
n, c, h, w = logits.shape
|
||
|
|
th, tw = target_hw
|
||
|
|
|
||
|
|
if (h, w) == (th, tw):
|
||
|
|
return logits
|
||
|
|
|
||
|
|
out = np.empty((n, c, th, tw), dtype=np.float32)
|
||
|
|
|
||
|
|
for bi in range(n):
|
||
|
|
for ci in range(c):
|
||
|
|
out[bi, ci] = cv2.resize(
|
||
|
|
logits[bi, ci].astype(np.float32),
|
||
|
|
(tw, th),
|
||
|
|
interpolation=cv2.INTER_LINEAR,
|
||
|
|
)
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def compute_mask_iou_between_preds(
|
||
|
|
pred_a: np.ndarray,
|
||
|
|
pred_b: np.ndarray,
|
||
|
|
num_classes: int,
|
||
|
|
) -> Tuple[List[Optional[float]], float, List[int]]:
|
||
|
|
"""
|
||
|
|
Mede IoU entre duas predições.
|
||
|
|
|
||
|
|
Classes ausentes nos dois mapas recebem None e NÃO entram no mIoU.
|
||
|
|
Isso evita o caso:
|
||
|
|
máscaras iguais, só classe 0 presente -> [1.0, None, None] -> mIoU=1.0
|
||
|
|
em vez de:
|
||
|
|
[1.0, 0.0, 0.0] -> mIoU=0.333
|
||
|
|
"""
|
||
|
|
a = pred_a.reshape(-1).astype(np.int64)
|
||
|
|
b = pred_b.reshape(-1).astype(np.int64)
|
||
|
|
|
||
|
|
valid = (a >= 0) & (a < num_classes) & (b >= 0) & (b < num_classes)
|
||
|
|
a = a[valid]
|
||
|
|
b = b[valid]
|
||
|
|
|
||
|
|
if a.size == 0:
|
||
|
|
return [None for _ in range(num_classes)], 0.0, []
|
||
|
|
|
||
|
|
cm = np.bincount(
|
||
|
|
num_classes * a + b,
|
||
|
|
minlength=num_classes * num_classes,
|
||
|
|
).reshape(num_classes, num_classes)
|
||
|
|
|
||
|
|
tp = np.diag(cm).astype(np.float64)
|
||
|
|
fp = cm.sum(axis=0).astype(np.float64) - tp
|
||
|
|
fn = cm.sum(axis=1).astype(np.float64) - tp
|
||
|
|
den = tp + fp + fn
|
||
|
|
|
||
|
|
iou_per_class: List[Optional[float]] = []
|
||
|
|
present_classes: List[int] = []
|
||
|
|
|
||
|
|
for cls in range(num_classes):
|
||
|
|
if den[cls] <= 0:
|
||
|
|
# Classe ausente nas duas predições.
|
||
|
|
iou_per_class.append(None)
|
||
|
|
else:
|
||
|
|
iou_per_class.append(float(tp[cls] / den[cls]))
|
||
|
|
present_classes.append(cls)
|
||
|
|
|
||
|
|
valid_ious = [x for x in iou_per_class if x is not None]
|
||
|
|
miou = float(np.mean(valid_ious)) if valid_ious else 0.0
|
||
|
|
|
||
|
|
return iou_per_class, miou, present_classes
|
||
|
|
|
||
|
|
|
||
|
|
def load_norm_stats(
|
||
|
|
path: Optional[Path],
|
||
|
|
channels: int,
|
||
|
|
channel_indices: List[int],
|
||
|
|
channel_names: List[str],
|
||
|
|
) -> Tuple[Optional[List[float]], Optional[List[float]], Optional[str]]:
|
||
|
|
if path is None or not path.is_file():
|
||
|
|
if path is not None:
|
||
|
|
print(f"[NORM] norm_stats não encontrado: {path}")
|
||
|
|
print("[NORM] Sem norm_stats. Usando tensor 0..1 sem padronização.")
|
||
|
|
return None, None, None
|
||
|
|
|
||
|
|
js = load_json(path)
|
||
|
|
mean = js.get("mean")
|
||
|
|
std = js.get("std")
|
||
|
|
names = js.get("channels", [])
|
||
|
|
|
||
|
|
if mean is None or std is None:
|
||
|
|
raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}")
|
||
|
|
|
||
|
|
max_idx = max(channel_indices)
|
||
|
|
|
||
|
|
if len(mean) <= max_idx or len(std) <= max_idx:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"norm_stats incompatível: precisa índices={channel_indices}, "
|
||
|
|
f"mean={len(mean)} std={len(std)}"
|
||
|
|
)
|
||
|
|
|
||
|
|
mean_sel = [float(mean[i]) for i in channel_indices]
|
||
|
|
std_sel = [float(std[i]) for i in channel_indices]
|
||
|
|
|
||
|
|
if names:
|
||
|
|
names_sel = [names[i] for i in channel_indices]
|
||
|
|
else:
|
||
|
|
names_sel = channel_names
|
||
|
|
|
||
|
|
print(f"[NORM] usando {path}")
|
||
|
|
print(f"[NORM] channels={names_sel}")
|
||
|
|
print(f"[NORM] mean={mean_sel}")
|
||
|
|
print(f"[NORM] std ={std_sel}")
|
||
|
|
|
||
|
|
return mean_sel, std_sel, str(path)
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_numpy_chw(chw: np.ndarray, mean: Optional[List[float]], std: Optional[List[float]]) -> np.ndarray:
|
||
|
|
if mean is None or std is None:
|
||
|
|
return chw.astype(np.float32)
|
||
|
|
|
||
|
|
mean_np = np.asarray(mean, dtype=np.float32).reshape(-1, 1, 1)
|
||
|
|
std_np = np.asarray(std, dtype=np.float32).reshape(-1, 1, 1)
|
||
|
|
std_np = np.clip(std_np, 1e-6, None)
|
||
|
|
|
||
|
|
return ((chw.astype(np.float32) - mean_np) / std_np).astype(np.float32)
|
||
|
|
|
||
|
|
|
||
|
|
def load_tensor(path: Path, channels: int, channel_indices: List[int]) -> np.ndarray:
|
||
|
|
arr = np.load(str(path)).astype(np.float32)
|
||
|
|
|
||
|
|
if arr.ndim != 3:
|
||
|
|
raise RuntimeError(f"Tensor inválido {path}: shape={arr.shape}, esperado 3D")
|
||
|
|
|
||
|
|
if arr.shape[0] in (3, 4, 5):
|
||
|
|
chw = arr
|
||
|
|
elif arr.shape[-1] in (3, 4, 5):
|
||
|
|
chw = np.transpose(arr, (2, 0, 1))
|
||
|
|
else:
|
||
|
|
raise RuntimeError(f"Tensor com layout inesperado: {path} shape={arr.shape}")
|
||
|
|
|
||
|
|
max_idx = max(channel_indices)
|
||
|
|
if chw.shape[0] <= max_idx:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Tensor {path} tem {chw.shape[0]} canais, "
|
||
|
|
f"mas precisa acessar índice {max_idx}."
|
||
|
|
)
|
||
|
|
|
||
|
|
chw = chw[channel_indices, :, :]
|
||
|
|
|
||
|
|
finite = np.isfinite(chw)
|
||
|
|
if finite.any():
|
||
|
|
mx = float(np.nanmax(chw[finite]))
|
||
|
|
if mx > 2.0 and mx <= 255.0:
|
||
|
|
chw = chw / 255.0
|
||
|
|
elif mx > 255.0:
|
||
|
|
chw = chw / 65535.0
|
||
|
|
|
||
|
|
chw = np.nan_to_num(chw, nan=0.0, posinf=1.0, neginf=0.0)
|
||
|
|
return np.clip(chw, 0.0, 1.0).astype(np.float32)
|
||
|
|
|
||
|
|
|
||
|
|
def collect_tensor_samples(root: Path, max_samples: int = 20, start_idx: int = 0) -> List[Path]:
|
||
|
|
tensor_paths: List[Path] = []
|
||
|
|
|
||
|
|
direct = root / "tensors"
|
||
|
|
if direct.is_dir():
|
||
|
|
tensor_paths.extend(sorted(direct.glob("*.npy")))
|
||
|
|
|
||
|
|
group_root = root / "group"
|
||
|
|
if group_root.is_dir():
|
||
|
|
for gdir in sorted(group_root.iterdir()):
|
||
|
|
tdir = gdir / "tensors"
|
||
|
|
if tdir.is_dir():
|
||
|
|
tensor_paths.extend(sorted(tdir.glob("*.npy")))
|
||
|
|
|
||
|
|
if not tensor_paths:
|
||
|
|
tensor_paths.extend(sorted(root.glob("**/tensors/*.npy")))
|
||
|
|
|
||
|
|
if not tensor_paths:
|
||
|
|
raise RuntimeError(f"Nenhum tensor .npy encontrado em: {root}")
|
||
|
|
|
||
|
|
start_idx = max(0, int(start_idx))
|
||
|
|
selected = tensor_paths[start_idx:]
|
||
|
|
|
||
|
|
if max_samples > 0:
|
||
|
|
selected = selected[:int(max_samples)]
|
||
|
|
|
||
|
|
return selected
|
||
|
|
|
||
|
|
|
||
|
|
def find_norm_stats(config: dict, config_dir: Path, save_dir: Path, explicit: Optional[str]) -> Optional[Path]:
|
||
|
|
if explicit:
|
||
|
|
return resolve_path(explicit, Path.cwd())
|
||
|
|
|
||
|
|
W, H = config.get("resolucao", [1024, 640])
|
||
|
|
dataset_path = config_dir / "dataset"
|
||
|
|
|
||
|
|
candidates = [
|
||
|
|
dataset_path / f"{int(W)}x{int(H)}" / "group" / "norm_stats.json",
|
||
|
|
save_dir / "norm_stats.json",
|
||
|
|
config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test") / config.get("stats_source_tag", "stacked_raw5") / "norm_stats.json",
|
||
|
|
]
|
||
|
|
|
||
|
|
for p in candidates:
|
||
|
|
if p.is_file():
|
||
|
|
return p
|
||
|
|
|
||
|
|
return candidates[0]
|
||
|
|
|
||
|
|
|
||
|
|
def nanmean_list(arr: np.ndarray) -> List[Optional[float]]:
|
||
|
|
if arr.size == 0:
|
||
|
|
return []
|
||
|
|
|
||
|
|
out = []
|
||
|
|
for col in range(arr.shape[1]):
|
||
|
|
v = arr[:, col]
|
||
|
|
v = v[~np.isnan(v)]
|
||
|
|
out.append(None if v.size == 0 else float(np.mean(v)))
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def nanmin_list(arr: np.ndarray) -> List[Optional[float]]:
|
||
|
|
if arr.size == 0:
|
||
|
|
return []
|
||
|
|
|
||
|
|
out = []
|
||
|
|
for col in range(arr.shape[1]):
|
||
|
|
v = arr[:, col]
|
||
|
|
v = v[~np.isnan(v)]
|
||
|
|
out.append(None if v.size == 0 else float(np.min(v)))
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_model_artifact_paths(
|
||
|
|
args,
|
||
|
|
config: dict,
|
||
|
|
config_dir: Path,
|
||
|
|
channels: int,
|
||
|
|
) -> Tuple[Path, Path, str]:
|
||
|
|
"""
|
||
|
|
Resolve checkpoint e ONNX.
|
||
|
|
|
||
|
|
Se --checkpoint ou --onnx forem informados, usa os caminhos informados.
|
||
|
|
Se ficarem vazios, monta a partir do config:
|
||
|
|
|
||
|
|
backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.pt
|
||
|
|
backup/{modelo}/{model_name}/{fusion_mode}_raw{channels}/{ckpt_name}.onnx
|
||
|
|
|
||
|
|
ckpt_name vem de:
|
||
|
|
config["ckpt_test"] ou "best_score"
|
||
|
|
"""
|
||
|
|
model = config.get("modelo", "segformer_b1")
|
||
|
|
model_name = config.get("model_name", "target_teached")
|
||
|
|
fusion_mode = config.get("fusion_mode", "stacked")
|
||
|
|
|
||
|
|
# Melhor usar o channels real computado pelo script,
|
||
|
|
# porque ele vem do input_channel_names já resolvido.
|
||
|
|
ch = int(channels)
|
||
|
|
|
||
|
|
ckpt_name = config.get("ckpt_test", "best_score")
|
||
|
|
|
||
|
|
base_dir = config_dir / "backup" / model / model_name / f"{fusion_mode}_raw{ch}"
|
||
|
|
|
||
|
|
if args.checkpoint:
|
||
|
|
checkpoint_path = resolve_path(args.checkpoint, Path.cwd())
|
||
|
|
else:
|
||
|
|
checkpoint_path = base_dir / f"{ckpt_name}.pt"
|
||
|
|
|
||
|
|
if args.onnx:
|
||
|
|
onnx_path = resolve_path(args.onnx, Path.cwd())
|
||
|
|
else:
|
||
|
|
onnx_path = base_dir / f"{ckpt_name}.onnx"
|
||
|
|
|
||
|
|
if checkpoint_path is None or not checkpoint_path.is_file():
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"Checkpoint não encontrado: {checkpoint_path}\n"
|
||
|
|
f"Dica: informe --checkpoint ou ajuste config['ckpt_test']."
|
||
|
|
)
|
||
|
|
|
||
|
|
if onnx_path is None or not onnx_path.is_file():
|
||
|
|
raise FileNotFoundError(
|
||
|
|
f"ONNX não encontrado: {onnx_path}\n"
|
||
|
|
f"Dica: informe --onnx ou ajuste config['ckpt_test']."
|
||
|
|
)
|
||
|
|
|
||
|
|
return checkpoint_path.resolve(), onnx_path.resolve(), str(ckpt_name)
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# PyTorch wrapper
|
||
|
|
# ============================================================
|
||
|
|
|
||
|
|
class MultiHeadTorchWrapper(nn.Module):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model: nn.Module,
|
||
|
|
output_heads: List[str],
|
||
|
|
resize_to_input: bool = False,
|
||
|
|
output_kind: str = "logits",
|
||
|
|
):
|
||
|
|
super().__init__()
|
||
|
|
self.model = model
|
||
|
|
self.output_heads = list(output_heads)
|
||
|
|
self.resize_to_input = bool(resize_to_input)
|
||
|
|
self.output_kind = str(output_kind).lower()
|
||
|
|
|
||
|
|
if self.output_kind not in ("logits", "mask"):
|
||
|
|
raise RuntimeError(f"output_kind inválido: {self.output_kind}")
|
||
|
|
|
||
|
|
def forward(self, pixel_values: torch.Tensor):
|
||
|
|
outputs: Dict[str, torch.Tensor] = self.model(pixel_values=pixel_values)
|
||
|
|
|
||
|
|
result = {}
|
||
|
|
input_hw = pixel_values.shape[-2:]
|
||
|
|
|
||
|
|
for head_name in self.output_heads:
|
||
|
|
logits = outputs[head_name]
|
||
|
|
|
||
|
|
if self.resize_to_input or self.output_kind == "mask":
|
||
|
|
if logits.shape[-2:] != input_hw:
|
||
|
|
logits = F.interpolate(
|
||
|
|
logits,
|
||
|
|
size=input_hw,
|
||
|
|
mode="bilinear",
|
||
|
|
align_corners=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
if self.output_kind == "mask":
|
||
|
|
result[head_name] = torch.argmax(logits, dim=1).to(torch.uint8)
|
||
|
|
else:
|
||
|
|
result[head_name] = logits
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# ONNX
|
||
|
|
# ============================================================
|
||
|
|
|
||
|
|
def create_onnx_session(
|
||
|
|
onnx_path: Path,
|
||
|
|
provider: str,
|
||
|
|
trt_home: Optional[str] = None,
|
||
|
|
trt_fp16: bool = True,
|
||
|
|
):
|
||
|
|
try:
|
||
|
|
import onnxruntime as ort
|
||
|
|
except ImportError:
|
||
|
|
raise ImportError(
|
||
|
|
"onnxruntime não está instalado. Instale com:\n"
|
||
|
|
" pip install onnxruntime-gpu\n"
|
||
|
|
"ou, para CPU:\n"
|
||
|
|
" pip install onnxruntime"
|
||
|
|
)
|
||
|
|
|
||
|
|
provider = provider.lower()
|
||
|
|
|
||
|
|
# ========================================================
|
||
|
|
# Windows/DLL helper para TensorRT
|
||
|
|
# ========================================================
|
||
|
|
if provider == "tensorrt":
|
||
|
|
trt_home = trt_home or os.environ.get("TRT_HOME", r"C:\dev\TensorRT-10.10.0.31")
|
||
|
|
|
||
|
|
dll_dirs = [
|
||
|
|
os.path.join(trt_home, "lib"),
|
||
|
|
os.path.join(trt_home, "bin"),
|
||
|
|
]
|
||
|
|
|
||
|
|
cuda_home = os.environ.get("CUDA_PATH")
|
||
|
|
if cuda_home:
|
||
|
|
dll_dirs.append(os.path.join(cuda_home, "bin"))
|
||
|
|
|
||
|
|
# Fallback comum para CUDA 12.4
|
||
|
|
dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin")
|
||
|
|
|
||
|
|
for dll_dir in dll_dirs:
|
||
|
|
if os.path.isdir(dll_dir):
|
||
|
|
try:
|
||
|
|
os.add_dll_directory(dll_dir)
|
||
|
|
print(f"[DLL] add_dll_directory: {dll_dir}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[DLL][WARN] falha em {dll_dir}: {e}")
|
||
|
|
|
||
|
|
available = ort.get_available_providers()
|
||
|
|
print(f"[ONNX] providers disponíveis: {available}")
|
||
|
|
|
||
|
|
sess_options = ort.SessionOptions()
|
||
|
|
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||
|
|
|
||
|
|
if provider == "cuda":
|
||
|
|
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
|
|
|
||
|
|
elif provider == "cpu":
|
||
|
|
providers = ["CPUExecutionProvider"]
|
||
|
|
|
||
|
|
elif provider == "tensorrt":
|
||
|
|
cache_dir = onnx_path.parent / "trt_cache"
|
||
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
trt_options = {
|
||
|
|
"device_id": 0,
|
||
|
|
"trt_fp16_enable": bool(trt_fp16),
|
||
|
|
|
||
|
|
# Cache para não reconstruir engine/timing toda vez.
|
||
|
|
"trt_engine_cache_enable": True,
|
||
|
|
"trt_engine_cache_path": str(cache_dir),
|
||
|
|
|
||
|
|
"trt_timing_cache_enable": True,
|
||
|
|
"trt_timing_cache_path": str(cache_dir),
|
||
|
|
|
||
|
|
# 4GB de workspace. Sua RTX 3070 lidou bem no benchmark.
|
||
|
|
"trt_max_workspace_size": 4 * 1024 * 1024 * 1024,
|
||
|
|
}
|
||
|
|
|
||
|
|
providers = [
|
||
|
|
("TensorrtExecutionProvider", trt_options),
|
||
|
|
"CUDAExecutionProvider",
|
||
|
|
"CPUExecutionProvider",
|
||
|
|
]
|
||
|
|
|
||
|
|
else:
|
||
|
|
raise RuntimeError(f"Provider desconhecido: {provider}")
|
||
|
|
|
||
|
|
requested_names = [
|
||
|
|
p[0] if isinstance(p, tuple) else p
|
||
|
|
for p in providers
|
||
|
|
]
|
||
|
|
|
||
|
|
providers_ok = [
|
||
|
|
p for p in providers
|
||
|
|
if (p[0] if isinstance(p, tuple) else p) in available
|
||
|
|
]
|
||
|
|
|
||
|
|
if not providers_ok:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Nenhum provider solicitado está disponível. "
|
||
|
|
f"Solicitado={requested_names}, disponível={available}"
|
||
|
|
)
|
||
|
|
|
||
|
|
sess = ort.InferenceSession(
|
||
|
|
str(onnx_path),
|
||
|
|
sess_options=sess_options,
|
||
|
|
providers=providers_ok,
|
||
|
|
)
|
||
|
|
|
||
|
|
active = sess.get_providers()
|
||
|
|
print(f"[ONNX] usando providers: {active}")
|
||
|
|
|
||
|
|
# Trava anti-burrice silenciosa: se pedir TensorRT, não pode cair para CPU/CUDA sem avisar.
|
||
|
|
if provider == "tensorrt" and "TensorrtExecutionProvider" not in active:
|
||
|
|
raise RuntimeError(
|
||
|
|
"TensorRTExecutionProvider foi solicitado, mas não ficou ativo. "
|
||
|
|
f"Providers ativos: {active}. "
|
||
|
|
"Provável causa: DLLs TensorRT fora do PATH/add_dll_directory, "
|
||
|
|
"versão incompatível ou fallback interno."
|
||
|
|
)
|
||
|
|
|
||
|
|
if provider == "cuda" and "CUDAExecutionProvider" not in active:
|
||
|
|
raise RuntimeError(
|
||
|
|
"CUDAExecutionProvider foi solicitado, mas não ficou ativo. "
|
||
|
|
f"Providers ativos: {active}."
|
||
|
|
)
|
||
|
|
|
||
|
|
return sess
|
||
|
|
|
||
|
|
|
||
|
|
def run_onnx(session, input_name: str, x_nchw: np.ndarray) -> Dict[str, np.ndarray]:
|
||
|
|
outputs = session.run(None, {input_name: x_nchw.astype(np.float32)})
|
||
|
|
output_names = [o.name for o in session.get_outputs()]
|
||
|
|
|
||
|
|
if len(outputs) != len(output_names):
|
||
|
|
raise RuntimeError("Quantidade de outputs ONNX inesperada.")
|
||
|
|
|
||
|
|
return {
|
||
|
|
name: arr.astype(np.float32)
|
||
|
|
for name, arr in zip(output_names, outputs)
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_onnx_output_names(onnx_outputs: Dict[str, np.ndarray], output_heads: List[str]) -> Dict[str, np.ndarray]:
|
||
|
|
"""
|
||
|
|
Converte:
|
||
|
|
semantic_logits -> semantic
|
||
|
|
vegetation_logits -> vegetation
|
||
|
|
etc.
|
||
|
|
"""
|
||
|
|
out = {}
|
||
|
|
|
||
|
|
for head in output_heads:
|
||
|
|
candidates = [
|
||
|
|
head,
|
||
|
|
f"{head}_logits",
|
||
|
|
f"{head}_mask",
|
||
|
|
f"output_{head}",
|
||
|
|
]
|
||
|
|
|
||
|
|
found = None
|
||
|
|
for c in candidates:
|
||
|
|
if c in onnx_outputs:
|
||
|
|
found = c
|
||
|
|
break
|
||
|
|
|
||
|
|
if found is None:
|
||
|
|
# fallback por ordem caso nomes estejam diferentes
|
||
|
|
keys = list(onnx_outputs.keys())
|
||
|
|
idx = output_heads.index(head)
|
||
|
|
if idx < len(keys):
|
||
|
|
found = keys[idx]
|
||
|
|
|
||
|
|
if found is None:
|
||
|
|
raise RuntimeError(f"Não encontrei saída ONNX para head={head}. Outputs={list(onnx_outputs.keys())}")
|
||
|
|
|
||
|
|
out[head] = onnx_outputs[found]
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
# ============================================================
|
||
|
|
# Main
|
||
|
|
# ============================================================
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
|
||
|
|
parser.add_argument("--config", default="config.json")
|
||
|
|
parser.add_argument("--checkpoint", default="")
|
||
|
|
parser.add_argument("--onnx", default="")
|
||
|
|
parser.add_argument("--train-script", default="_8_train_multihead.py")
|
||
|
|
parser.add_argument("--labelmap", default="dataset/labelmap.txt")
|
||
|
|
|
||
|
|
parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
|
||
|
|
parser.add_argument("--root_override", default=None)
|
||
|
|
|
||
|
|
parser.add_argument("--norm_stats", default=None)
|
||
|
|
parser.add_argument("--max_samples", type=int, default=20)
|
||
|
|
parser.add_argument("--start_idx", type=int, default=0)
|
||
|
|
|
||
|
|
parser.add_argument("--device", default="cuda", choices=["cuda", "cpu"])
|
||
|
|
parser.add_argument("--onnx_provider", default="cuda", choices=["cuda", "cpu", "tensorrt"])
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--trt_home",
|
||
|
|
default="C:\\dev\\TensorRT-10.10.0.31",
|
||
|
|
help="Pasta raiz do TensorRT. Ex: C:\\dev\\TensorRT-10.10.0.31. Se omitido, usa TRT_HOME ou fallback padrão.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--trt_no_fp16",
|
||
|
|
action="store_true",
|
||
|
|
help="Desativa FP16 no TensorRT. Normalmente NÃO usar; deixamos FP16 ligado.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--torch_no_amp",
|
||
|
|
action="store_true",
|
||
|
|
help="Desativa AMP no PyTorch. Recomendado para comparação mais rígida contra ONNX FP32.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--resize_torch_to_input",
|
||
|
|
action="store_true",
|
||
|
|
help="Força saída PyTorch redimensionada para HxW antes de comparar.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--onnx_has_norm",
|
||
|
|
action="store_true",
|
||
|
|
help="Use quando o ONNX já inclui normalização interna. Nesse caso o ONNX recebe tensor 0..1, não tensor normalizado.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--onnx_output_kind",
|
||
|
|
default="logits",
|
||
|
|
choices=["logits", "mask"],
|
||
|
|
help="Tipo de saída do ONNX: logits para modelo cru, mask para ONNX com argmax/postprocess embutido.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--compare_at_input_size",
|
||
|
|
action="store_true",
|
||
|
|
help="Redimensiona ambos os outputs para HxW antes de comparar.",
|
||
|
|
)
|
||
|
|
|
||
|
|
parser.add_argument(
|
||
|
|
"--save_report",
|
||
|
|
default=None,
|
||
|
|
help="Caminho do JSON de relatório. Se omitido, salva ao lado do ONNX.",
|
||
|
|
)
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
config_path = resolve_path(args.config, Path.cwd())
|
||
|
|
train_script_path = resolve_path(args.train_script, Path.cwd())
|
||
|
|
labelmap_path = resolve_path(args.labelmap, Path.cwd())
|
||
|
|
|
||
|
|
if config_path is None or not config_path.is_file():
|
||
|
|
raise FileNotFoundError(f"Config não encontrado: {config_path}")
|
||
|
|
if train_script_path is None or not train_script_path.is_file():
|
||
|
|
raise FileNotFoundError(f"Train script não encontrado: {train_script_path}")
|
||
|
|
if labelmap_path is None or not labelmap_path.is_file():
|
||
|
|
raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")
|
||
|
|
|
||
|
|
config_dir = config_path.parent
|
||
|
|
config = load_json(config_path)
|
||
|
|
|
||
|
|
train_mod = import_train_module(train_script_path)
|
||
|
|
|
||
|
|
W, H = config.get("resolucao", [1024, 640])
|
||
|
|
W = int(W)
|
||
|
|
H = int(H)
|
||
|
|
|
||
|
|
backbone = config.get("backbone", "nvidia/mit-b1")
|
||
|
|
input_channel_names = train_mod.get_input_channel_names(config)
|
||
|
|
input_channel_indices = train_mod.get_input_channel_indices(config)
|
||
|
|
channels = len(input_channel_names)
|
||
|
|
|
||
|
|
checkpoint_path, onnx_path, ckpt_name = resolve_model_artifact_paths(
|
||
|
|
args=args,
|
||
|
|
config=config,
|
||
|
|
config_dir=config_dir,
|
||
|
|
channels=channels,
|
||
|
|
)
|
||
|
|
|
||
|
|
semantic_id2label, semantic_label2id, ignore_from_labelmap = train_mod.load_labelmap(
|
||
|
|
str(labelmap_path)
|
||
|
|
)
|
||
|
|
|
||
|
|
heads_config = train_mod.build_heads_config(
|
||
|
|
config,
|
||
|
|
ignore_index=int(ignore_from_labelmap)
|
||
|
|
)
|
||
|
|
|
||
|
|
heads_config["semantic"]["num_classes"] = int(len(semantic_id2label))
|
||
|
|
heads_config["semantic"]["ignore_index"] = int(ignore_from_labelmap)
|
||
|
|
|
||
|
|
output_heads = list(heads_config.keys())
|
||
|
|
|
||
|
|
save_dir = (
|
||
|
|
config_dir
|
||
|
|
/ "backup"
|
||
|
|
/ config.get("modelo", "segformer_b1")
|
||
|
|
/ config.get("model_name", "test")
|
||
|
|
/ f"{config.get('fusion_mode', 'stacked')}_raw{channels}"
|
||
|
|
)
|
||
|
|
|
||
|
|
norm_stats_path = find_norm_stats(
|
||
|
|
config=config,
|
||
|
|
config_dir=config_dir,
|
||
|
|
save_dir=save_dir,
|
||
|
|
explicit=args.norm_stats,
|
||
|
|
)
|
||
|
|
|
||
|
|
mean, std, norm_stats_used = load_norm_stats(
|
||
|
|
norm_stats_path,
|
||
|
|
channels=channels,
|
||
|
|
channel_indices=input_channel_indices,
|
||
|
|
channel_names=input_channel_names,
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.root_override:
|
||
|
|
root = resolve_path(args.root_override, Path.cwd())
|
||
|
|
else:
|
||
|
|
root = (config_dir / "dataset" / "split" / args.split_folder).resolve()
|
||
|
|
|
||
|
|
if root is None or not root.is_dir():
|
||
|
|
raise FileNotFoundError(f"Root de dados não encontrado: {root}")
|
||
|
|
|
||
|
|
samples = collect_tensor_samples(
|
||
|
|
root=root,
|
||
|
|
max_samples=args.max_samples,
|
||
|
|
start_idx=args.start_idx,
|
||
|
|
)
|
||
|
|
|
||
|
|
use_cuda = args.device == "cuda" and torch.cuda.is_available()
|
||
|
|
device = torch.device("cuda" if use_cuda else "cpu")
|
||
|
|
|
||
|
|
if args.device == "cuda" and not torch.cuda.is_available():
|
||
|
|
print("[WARN] CUDA indisponível. Usando CPU no PyTorch.")
|
||
|
|
|
||
|
|
print("==========================================")
|
||
|
|
print("Validate PyTorch vs ONNX")
|
||
|
|
print(f"Config : {config_path}")
|
||
|
|
print(f"Checkpoint : {checkpoint_path}")
|
||
|
|
print(f"ONNX : {onnx_path}")
|
||
|
|
print(f"Root : {root}")
|
||
|
|
print(f"Samples : {len(samples)}")
|
||
|
|
print(f"Backbone : {backbone}")
|
||
|
|
print(f"Input shape : [1, {channels}, {H}, {W}]")
|
||
|
|
print(f"Channels : {input_channel_names} idx={input_channel_indices}")
|
||
|
|
print(f"Heads : {output_heads}")
|
||
|
|
print(f"Device : {device}")
|
||
|
|
print(f"ONNX provider: {args.onnx_provider}")
|
||
|
|
print(f"Torch AMP : {not args.torch_no_amp and device.type == 'cuda'}")
|
||
|
|
print(f"Compare HxW : {args.compare_at_input_size}")
|
||
|
|
print("==========================================")
|
||
|
|
|
||
|
|
print("[MODEL] Montando PyTorch...")
|
||
|
|
model = train_mod.build_model(
|
||
|
|
backbone=backbone,
|
||
|
|
channels=channels,
|
||
|
|
heads_config=heads_config,
|
||
|
|
semantic_id2label=semantic_id2label,
|
||
|
|
semantic_label2id=semantic_label2id,
|
||
|
|
)
|
||
|
|
|
||
|
|
ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
|
||
|
|
if "model" not in ckpt:
|
||
|
|
raise RuntimeError("Checkpoint não contém chave 'model'.")
|
||
|
|
|
||
|
|
model.load_state_dict(ckpt["model"], strict=True)
|
||
|
|
model.to(device)
|
||
|
|
model.eval()
|
||
|
|
|
||
|
|
torch_wrapper = MultiHeadTorchWrapper(
|
||
|
|
model=model,
|
||
|
|
output_heads=output_heads,
|
||
|
|
resize_to_input=args.resize_torch_to_input or args.onnx_output_kind == "mask",
|
||
|
|
output_kind=args.onnx_output_kind,
|
||
|
|
).to(device)
|
||
|
|
torch_wrapper.eval()
|
||
|
|
|
||
|
|
print("[ONNX] Carregando sessão...")
|
||
|
|
onnx_session = create_onnx_session(
|
||
|
|
onnx_path=onnx_path,
|
||
|
|
provider=args.onnx_provider,
|
||
|
|
trt_home=args.trt_home,
|
||
|
|
trt_fp16=not args.trt_no_fp16,
|
||
|
|
)
|
||
|
|
onnx_input_name = onnx_session.get_inputs()[0].name
|
||
|
|
print(f"[ONNX] input name: {onnx_input_name}")
|
||
|
|
print(f"[ONNX] outputs: {[o.name for o in onnx_session.get_outputs()]}")
|
||
|
|
|
||
|
|
per_head_accum = {
|
||
|
|
h: {
|
||
|
|
"n": 0,
|
||
|
|
"logits_abs_mean": [],
|
||
|
|
"logits_abs_max": [],
|
||
|
|
"prob_abs_mean": [],
|
||
|
|
"prob_abs_max": [],
|
||
|
|
"argmax_equal_ratio": [],
|
||
|
|
"pred_miou_torch_vs_onnx": [],
|
||
|
|
"pred_iou_per_class": [],
|
||
|
|
}
|
||
|
|
for h in output_heads
|
||
|
|
}
|
||
|
|
|
||
|
|
sample_reports = []
|
||
|
|
|
||
|
|
for i, tensor_path in enumerate(samples):
|
||
|
|
chw01 = load_tensor(
|
||
|
|
tensor_path,
|
||
|
|
channels=channels,
|
||
|
|
channel_indices=input_channel_indices,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Garante resolução do contrato.
|
||
|
|
if chw01.shape[-2:] != (H, W):
|
||
|
|
hwc = np.transpose(chw01, (1, 2, 0))
|
||
|
|
hwc = cv2.resize(hwc, (W, H), interpolation=cv2.INTER_LINEAR)
|
||
|
|
chw01 = np.transpose(hwc, (2, 0, 1)).astype(np.float32)
|
||
|
|
|
||
|
|
chw_norm = normalize_numpy_chw(chw01, mean=mean, std=std)
|
||
|
|
|
||
|
|
# PyTorch continua recebendo normalizado, porque o modelo PyTorch puro espera isso.
|
||
|
|
x_torch_np = np.expand_dims(chw_norm, axis=0).astype(np.float32)
|
||
|
|
|
||
|
|
# ONNX novo com include_norm recebe 0..1 cru.
|
||
|
|
if args.onnx_has_norm:
|
||
|
|
x_onnx_np = np.expand_dims(chw01, axis=0).astype(np.float32)
|
||
|
|
else:
|
||
|
|
x_onnx_np = x_torch_np
|
||
|
|
|
||
|
|
x_torch = torch.from_numpy(x_torch_np).to(device, non_blocking=True)
|
||
|
|
|
||
|
|
if device.type == "cuda":
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
|
||
|
|
with torch.inference_mode():
|
||
|
|
with torch.autocast(
|
||
|
|
device_type="cuda",
|
||
|
|
dtype=torch.float16,
|
||
|
|
enabled=(not args.torch_no_amp and device.type == "cuda"),
|
||
|
|
):
|
||
|
|
torch_outputs_t = torch_wrapper(x_torch)
|
||
|
|
|
||
|
|
if device.type == "cuda":
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
torch_ms = (time.perf_counter() - t0) * 1000.0
|
||
|
|
|
||
|
|
torch_outputs = {
|
||
|
|
h: v.detach().float().cpu().numpy()
|
||
|
|
for h, v in torch_outputs_t.items()
|
||
|
|
}
|
||
|
|
|
||
|
|
t0 = time.perf_counter()
|
||
|
|
onnx_raw_outputs = run_onnx(onnx_session, onnx_input_name, x_onnx_np)
|
||
|
|
onnx_ms = (time.perf_counter() - t0) * 1000.0
|
||
|
|
|
||
|
|
onnx_outputs = normalize_onnx_output_names(
|
||
|
|
onnx_raw_outputs,
|
||
|
|
output_heads=output_heads,
|
||
|
|
)
|
||
|
|
|
||
|
|
report_item = {
|
||
|
|
"idx": i,
|
||
|
|
"tensor": str(tensor_path),
|
||
|
|
"torch_ms": float(torch_ms),
|
||
|
|
"onnx_ms": float(onnx_ms),
|
||
|
|
"heads": {},
|
||
|
|
}
|
||
|
|
|
||
|
|
print(f"\n[{i + 1:03d}/{len(samples):03d}] {tensor_path.name} | torch={torch_ms:.2f}ms | onnx={onnx_ms:.2f}ms")
|
||
|
|
|
||
|
|
for head in output_heads:
|
||
|
|
pt = torch_outputs[head]
|
||
|
|
ox = onnx_outputs[head]
|
||
|
|
|
||
|
|
if args.onnx_output_kind == "mask":
|
||
|
|
# Esperado:
|
||
|
|
# PT : [1,H,W]
|
||
|
|
# ONNX : [1,H,W]
|
||
|
|
pt_mask = np.asarray(pt).astype(np.uint8)
|
||
|
|
ox_mask = np.asarray(ox).astype(np.uint8)
|
||
|
|
|
||
|
|
if pt_mask.ndim == 3:
|
||
|
|
pt_mask = pt_mask[0]
|
||
|
|
if ox_mask.ndim == 3:
|
||
|
|
ox_mask = ox_mask[0]
|
||
|
|
|
||
|
|
if pt_mask.shape != ox_mask.shape:
|
||
|
|
ox_mask = cv2.resize(
|
||
|
|
ox_mask,
|
||
|
|
(pt_mask.shape[1], pt_mask.shape[0]),
|
||
|
|
interpolation=cv2.INTER_NEAREST,
|
||
|
|
)
|
||
|
|
|
||
|
|
equal_ratio = float(np.mean(pt_mask == ox_mask))
|
||
|
|
|
||
|
|
num_classes = int(heads_config[head]["num_classes"])
|
||
|
|
iou_per_class, miou, present_classes = compute_mask_iou_between_preds(
|
||
|
|
pt_mask,
|
||
|
|
ox_mask,
|
||
|
|
num_classes=num_classes,
|
||
|
|
)
|
||
|
|
|
||
|
|
head_report = {
|
||
|
|
"torch_shape": list(np.asarray(pt).shape),
|
||
|
|
"onnx_shape": list(np.asarray(ox).shape),
|
||
|
|
"compare_shape": list(pt_mask.shape),
|
||
|
|
"logits_abs_mean": None,
|
||
|
|
"logits_abs_max": None,
|
||
|
|
"prob_abs_mean": None,
|
||
|
|
"prob_abs_max": None,
|
||
|
|
"argmax_equal_ratio": float(equal_ratio),
|
||
|
|
"pred_miou_torch_vs_onnx": float(miou),
|
||
|
|
"pred_iou_per_class": [
|
||
|
|
None if x is None else float(x)
|
||
|
|
for x in iou_per_class
|
||
|
|
],
|
||
|
|
"pred_present_classes": [int(x) for x in present_classes],
|
||
|
|
}
|
||
|
|
|
||
|
|
print(
|
||
|
|
f" {head:<10} "
|
||
|
|
f"shape PT={tuple(np.asarray(pt).shape)} ONNX={tuple(np.asarray(ox).shape)} "
|
||
|
|
f"CMP={tuple(pt_mask.shape)} | "
|
||
|
|
f"mask_equal={equal_ratio * 100:.3f}% "
|
||
|
|
f"mIoU={miou:.6f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
else:
|
||
|
|
pt = np.asarray(pt).astype(np.float32)
|
||
|
|
ox = np.asarray(ox).astype(np.float32)
|
||
|
|
|
||
|
|
compare_hw = (H, W) if args.compare_at_input_size else None
|
||
|
|
|
||
|
|
if pt.shape != ox.shape:
|
||
|
|
compare_hw = (H, W)
|
||
|
|
|
||
|
|
if compare_hw is not None:
|
||
|
|
pt_cmp = resize_logits_np_nchw(pt, compare_hw)
|
||
|
|
ox_cmp = resize_logits_np_nchw(ox, compare_hw)
|
||
|
|
else:
|
||
|
|
pt_cmp = pt
|
||
|
|
ox_cmp = ox
|
||
|
|
|
||
|
|
if pt_cmp.shape != ox_cmp.shape:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"Shape incompatível na head {head}: "
|
||
|
|
f"torch={pt_cmp.shape}, onnx={ox_cmp.shape}"
|
||
|
|
)
|
||
|
|
|
||
|
|
diff_logits = np.abs(pt_cmp - ox_cmp)
|
||
|
|
|
||
|
|
prob_pt = softmax_np(pt_cmp, axis=1)
|
||
|
|
prob_ox = softmax_np(ox_cmp, axis=1)
|
||
|
|
diff_prob = np.abs(prob_pt - prob_ox)
|
||
|
|
|
||
|
|
pred_pt = np.argmax(prob_pt, axis=1)[0].astype(np.uint8)
|
||
|
|
pred_ox = np.argmax(prob_ox, axis=1)[0].astype(np.uint8)
|
||
|
|
|
||
|
|
equal_ratio = float(np.mean(pred_pt == pred_ox))
|
||
|
|
|
||
|
|
num_classes = int(heads_config[head]["num_classes"])
|
||
|
|
iou_per_class, miou, present_classes = compute_mask_iou_between_preds(
|
||
|
|
pred_pt,
|
||
|
|
pred_ox,
|
||
|
|
num_classes=num_classes,
|
||
|
|
)
|
||
|
|
|
||
|
|
head_report = {
|
||
|
|
"torch_shape": list(pt.shape),
|
||
|
|
"onnx_shape": list(ox.shape),
|
||
|
|
"compare_shape": list(pt_cmp.shape),
|
||
|
|
"logits_abs_mean": float(diff_logits.mean()),
|
||
|
|
"logits_abs_max": float(diff_logits.max()),
|
||
|
|
"prob_abs_mean": float(diff_prob.mean()),
|
||
|
|
"prob_abs_max": float(diff_prob.max()),
|
||
|
|
"argmax_equal_ratio": float(equal_ratio),
|
||
|
|
"pred_miou_torch_vs_onnx": float(miou),
|
||
|
|
"pred_iou_per_class": [
|
||
|
|
None if x is None else float(x)
|
||
|
|
for x in iou_per_class
|
||
|
|
],
|
||
|
|
"pred_present_classes": [int(x) for x in present_classes],
|
||
|
|
}
|
||
|
|
|
||
|
|
print(
|
||
|
|
f" {head:<10} "
|
||
|
|
f"shape PT={tuple(pt.shape)} ONNX={tuple(ox.shape)} CMP={tuple(pt_cmp.shape)} | "
|
||
|
|
f"logit_mean={head_report['logits_abs_mean']:.6g} "
|
||
|
|
f"prob_mean={head_report['prob_abs_mean']:.6g} "
|
||
|
|
f"argmax_equal={equal_ratio * 100:.3f}% "
|
||
|
|
f"mIoU={miou:.6f}"
|
||
|
|
)
|
||
|
|
|
||
|
|
report_item["heads"][head] = head_report
|
||
|
|
|
||
|
|
acc = per_head_accum[head]
|
||
|
|
acc["n"] += 1
|
||
|
|
|
||
|
|
if head_report["logits_abs_mean"] is not None:
|
||
|
|
acc["logits_abs_mean"].append(head_report["logits_abs_mean"])
|
||
|
|
|
||
|
|
if head_report["logits_abs_max"] is not None:
|
||
|
|
acc["logits_abs_max"].append(head_report["logits_abs_max"])
|
||
|
|
|
||
|
|
if head_report["prob_abs_mean"] is not None:
|
||
|
|
acc["prob_abs_mean"].append(head_report["prob_abs_mean"])
|
||
|
|
|
||
|
|
if head_report["prob_abs_max"] is not None:
|
||
|
|
acc["prob_abs_max"].append(head_report["prob_abs_max"])
|
||
|
|
|
||
|
|
acc["argmax_equal_ratio"].append(head_report["argmax_equal_ratio"])
|
||
|
|
acc["pred_miou_torch_vs_onnx"].append(head_report["pred_miou_torch_vs_onnx"])
|
||
|
|
acc["pred_iou_per_class"].append(head_report["pred_iou_per_class"])
|
||
|
|
|
||
|
|
#print(
|
||
|
|
# f" {head:<10} "
|
||
|
|
# f"shape PT={tuple(pt.shape)} ONNX={tuple(ox.shape)} CMP={tuple(pt_cmp.shape)} | "
|
||
|
|
# f"logit_mean={head_report['logits_abs_mean']:.6g} "
|
||
|
|
# f"prob_mean={head_report['prob_abs_mean']:.6g} "
|
||
|
|
# f"argmax_equal={equal_ratio * 100:.3f}% "
|
||
|
|
# f"mIoU={miou:.6f}"
|
||
|
|
#)
|
||
|
|
|
||
|
|
sample_reports.append(report_item)
|
||
|
|
|
||
|
|
summary = {
|
||
|
|
"config": str(config_path),
|
||
|
|
"checkpoint": str(checkpoint_path),
|
||
|
|
"ckpt_name": ckpt_name,
|
||
|
|
"onnx": str(onnx_path),
|
||
|
|
"root": str(root),
|
||
|
|
"samples": len(samples),
|
||
|
|
"input_shape": [1, channels, H, W],
|
||
|
|
"input_channel_names": input_channel_names,
|
||
|
|
"input_channel_indices": input_channel_indices,
|
||
|
|
"heads": output_heads,
|
||
|
|
"norm_stats_used": norm_stats_used,
|
||
|
|
"torch_amp": bool(not args.torch_no_amp and device.type == "cuda"),
|
||
|
|
"onnx_provider": args.onnx_provider,
|
||
|
|
"trt_home": args.trt_home or os.environ.get("TRT_HOME", None),
|
||
|
|
"trt_fp16": bool(not args.trt_no_fp16),
|
||
|
|
"compare_at_input_size": bool(args.compare_at_input_size),
|
||
|
|
"per_head": {},
|
||
|
|
"sample_reports": sample_reports,
|
||
|
|
}
|
||
|
|
|
||
|
|
print("\n========== RESUMO ==========")
|
||
|
|
|
||
|
|
for head, acc in per_head_accum.items():
|
||
|
|
if acc["n"] <= 0:
|
||
|
|
continue
|
||
|
|
|
||
|
|
ious_raw = acc["pred_iou_per_class"]
|
||
|
|
|
||
|
|
if ious_raw:
|
||
|
|
ious_arr = np.array(
|
||
|
|
[
|
||
|
|
[np.nan if x is None else float(x) for x in row]
|
||
|
|
for row in ious_raw
|
||
|
|
],
|
||
|
|
dtype=np.float64,
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
ious_arr = np.empty((0, 0), dtype=np.float64)
|
||
|
|
|
||
|
|
def safe_mean(values):
|
||
|
|
return None if not values else float(np.mean(values))
|
||
|
|
|
||
|
|
def safe_max(values):
|
||
|
|
return None if not values else float(np.max(values))
|
||
|
|
|
||
|
|
head_summary = {
|
||
|
|
"n": int(acc["n"]),
|
||
|
|
"logits_abs_mean_avg": safe_mean(acc["logits_abs_mean"]),
|
||
|
|
"logits_abs_mean_max": safe_max(acc["logits_abs_mean"]),
|
||
|
|
"logits_abs_max_avg": safe_mean(acc["logits_abs_max"]),
|
||
|
|
"logits_abs_max_max": safe_max(acc["logits_abs_max"]),
|
||
|
|
"prob_abs_mean_avg": safe_mean(acc["prob_abs_mean"]),
|
||
|
|
"prob_abs_mean_max": safe_max(acc["prob_abs_mean"]),
|
||
|
|
"prob_abs_max_avg": safe_mean(acc["prob_abs_max"]),
|
||
|
|
"prob_abs_max_max": safe_max(acc["prob_abs_max"]),
|
||
|
|
"argmax_equal_ratio_avg": float(np.mean(acc["argmax_equal_ratio"])),
|
||
|
|
"argmax_equal_ratio_min": float(np.min(acc["argmax_equal_ratio"])),
|
||
|
|
"pred_miou_avg": float(np.mean(acc["pred_miou_torch_vs_onnx"])),
|
||
|
|
"pred_miou_min": float(np.min(acc["pred_miou_torch_vs_onnx"])),
|
||
|
|
"pred_iou_per_class_avg": nanmean_list(ious_arr),
|
||
|
|
"pred_iou_per_class_min": nanmin_list(ious_arr),
|
||
|
|
}
|
||
|
|
|
||
|
|
summary["per_head"][head] = head_summary
|
||
|
|
|
||
|
|
def fmt_iou_list(values):
|
||
|
|
return [
|
||
|
|
None if x is None else round(float(x), 6)
|
||
|
|
for x in values
|
||
|
|
]
|
||
|
|
|
||
|
|
def fmt_optional(v, casas=8):
|
||
|
|
if v is None:
|
||
|
|
return "N/A"
|
||
|
|
return f"{float(v):.{casas}f}"
|
||
|
|
|
||
|
|
print(f"\n[{head}]")
|
||
|
|
print(f" logits_abs_mean avg : {fmt_optional(head_summary['logits_abs_mean_avg'])}")
|
||
|
|
print(f" prob_abs_mean avg : {fmt_optional(head_summary['prob_abs_mean_avg'])}")
|
||
|
|
print(f" argmax_equal avg : {head_summary['argmax_equal_ratio_avg'] * 100:.4f}%")
|
||
|
|
print(f" argmax_equal min : {head_summary['argmax_equal_ratio_min'] * 100:.4f}%")
|
||
|
|
print(f" pred_mIoU avg : {head_summary['pred_miou_avg']:.8f}")
|
||
|
|
print(f" pred_mIoU min : {head_summary['pred_miou_min']:.8f}")
|
||
|
|
print(f" IoU/classes avg : {fmt_iou_list(head_summary['pred_iou_per_class_avg'])}")
|
||
|
|
print(f" IoU/classes min : {fmt_iou_list(head_summary['pred_iou_per_class_min'])}")
|
||
|
|
|
||
|
|
if args.save_report:
|
||
|
|
report_path = resolve_path(args.save_report, Path.cwd())
|
||
|
|
else:
|
||
|
|
suffix = f".validate_{args.onnx_provider}_report.json"
|
||
|
|
report_path = onnx_path.with_suffix(suffix)
|
||
|
|
|
||
|
|
save_json(report_path, summary)
|
||
|
|
print(f"\n[OK] Relatório salvo em: {report_path}")
|
||
|
|
|
||
|
|
print("\nValidação finalizada.")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|