Zendion/agrobot_base
Zendion
/
agrobot_base
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
agrobot_base/Python/OAK/datasets/_12_benchmark_onnx.py

1250 lines
39 KiB
Python
Raw Normal View History

Implementação ONNX/TensorRT
2026-05-22 22:32:11 +00:00
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
_12_benchmark_visual_onnx.py
Benchmark PyTorch vs ONNX Runtime para o Visual Worker.
Contrato esperado:
- Modelo: SegFormer + cabeça auxiliar LabelHead
- Entrada: RGB [N, 3, H, W], float32
- ONNX exportado pelo _10_export_visual_onnx.py
- Saídas ONNX típicas:
semantic_logits: [N, 2, H, W]
label_probs : [N, 8]
Mede:
- PyTorch FP32
- PyTorch AMP/FP16
- ONNX Runtime CUDA / CPU / TensorRT
Importante:
- Este benchmark mede só inferência do modelo.
- As imagens são carregadas e pré-processadas antes da medição.
- Para ONNX com --include-norm, use --onnx_has_norm.
Exemplos:
# Benchmark completo CUDA
python _12_benchmark_visual_onnx.py ^
--config config.json ^
--onnx_provider cuda ^
--onnx_has_norm ^
--max_samples 50 ^
--warmup 10 ^
--repeat 5
# Benchmark TensorRT, medindo só ONNX TensorRT
python _12_benchmark_visual_onnx.py ^
--config config.json ^
--onnx_provider tensorrt ^
--onnx_has_norm ^
--max_samples 50 ^
--warmup 20 ^
--repeat 10 ^
--skip_torch_fp32 ^
--skip_torch_amp
# Pasta externa com imagens soltas
python _12_benchmark_visual_onnx.py ^
--config config.json ^
--test_folder .\oak-d\dataset\split\train\group\naonavegavel_navegavel\images\ ^
--onnx_provider tensorrt ^
--onnx_has_norm ^
--warmup 20 ^
--repeat 10
"""
from __future__ import annotations
import os
import gc
import csv
import glob
import json
import time
import argparse
from dataclasses import dataclass
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
from transformers import SegformerForSemanticSegmentation
# ============================================================
# Utils gerais
# ============================================================
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 save_csv(path: str | Path, rows: List[dict]):
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
if not rows:
return
keys = list(rows[0].keys())
with path.open("w", newline="", encoding="utf-8") as f:
w = csv.DictWriter(f, fieldnames=keys)
w.writeheader()
w.writerows(rows)
def resolve_path(path_like: Optional[str], base: Optional[Path] = None) -> Optional[Path]:
if path_like is None or str(path_like).strip() == "":
return None
p = Path(path_like)
if p.is_absolute():
return p
if base is None:
base = Path.cwd()
return (base / p).resolve()
def synchronize_if_cuda(device: torch.device):
if device.type == "cuda":
torch.cuda.synchronize()
def clear_cuda():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.synchronize()
def summarize_times(times_ms: List[float]) -> dict:
arr = np.asarray(times_ms, dtype=np.float64)
if arr.size == 0:
return {
"n": 0,
"mean_ms": 0.0,
"median_ms": 0.0,
"min_ms": 0.0,
"max_ms": 0.0,
"p90_ms": 0.0,
"p95_ms": 0.0,
"p99_ms": 0.0,
"fps_mean": 0.0,
"fps_median": 0.0,
"fps_p95_latency": 0.0,
"fps_p99_latency": 0.0,
}
mean_ms = float(arr.mean())
median_ms = float(np.median(arr))
p95_ms = float(np.percentile(arr, 95))
p99_ms = float(np.percentile(arr, 99))
return {
"n": int(arr.size),
"mean_ms": mean_ms,
"median_ms": median_ms,
"min_ms": float(arr.min()),
"max_ms": float(arr.max()),
"p90_ms": float(np.percentile(arr, 90)),
"p95_ms": p95_ms,
"p99_ms": p99_ms,
"fps_mean": float(1000.0 / mean_ms) if mean_ms > 0 else 0.0,
"fps_median": float(1000.0 / median_ms) if median_ms > 0 else 0.0,
"fps_p95_latency": float(1000.0 / p95_ms) if p95_ms > 0 else 0.0,
"fps_p99_latency": float(1000.0 / p99_ms) if p99_ms > 0 else 0.0,
}
# ============================================================
# Labelmap / paths / config
# ============================================================
def _try_int(text: str) -> Optional[int]:
try:
return int(str(text).strip())
except Exception:
return None
def clean_label_name(raw_name: str) -> str:
name = str(raw_name).strip()
if "::" in name:
name = name.split("::", 1)[0].strip()
if ":" in name:
left, right = name.split(":", 1)
right_clean = right.replace(",", "").replace(" ", "")
if right_clean.isdigit():
name = left.strip()
return name.strip()
def load_labelmap(labelmap_path: Path) -> Tuple[Dict[int, str], Dict[str, int], int]:
if not labelmap_path.is_file():
raise FileNotFoundError(f"Labelmap não encontrado: {labelmap_path}")
id2label: Dict[int, str] = {}
ignore_index = 255
next_id = 0
with labelmap_path.open("r", encoding="utf-8") as f:
for raw_line in f:
line = raw_line.strip()
if not line or line.startswith("#"):
continue
lower = line.lower()
if lower.startswith("ignore") or lower.startswith("ignore_index"):
for sep in ("=", ":", ",", " "):
if sep in line:
maybe = _try_int(line.split(sep)[-1])
if maybe is not None:
ignore_index = maybe
break
continue
cls_id: Optional[int] = None
cls_name: Optional[str] = None
if "::" in line and ":" in line:
before = line.split("::", 1)[0].strip()
maybe_name = before.split(":", 1)[0].strip()
if maybe_name:
cls_id = next_id
cls_name = maybe_name
if cls_id is None:
for sep in (":", ",", "\t", " "):
if sep in line:
parts = [p.strip() for p in line.split(sep) if p.strip()]
if len(parts) >= 2:
left_id = _try_int(parts[0])
right_id = _try_int(parts[-1])
if left_id is not None:
cls_id = left_id
cls_name = sep.join(parts[1:]).strip() if sep in (":", ",") else " ".join(parts[1:]).strip()
break
if right_id is not None:
cls_id = right_id
cls_name = sep.join(parts[:-1]).strip() if sep in (":", ",") else " ".join(parts[:-1]).strip()
break
if cls_id is None:
cls_id = next_id
cls_name = line
if cls_name is None or cls_name == "":
raise RuntimeError(f"Linha inválida no labelmap: {raw_line!r}")
id2label[int(cls_id)] = clean_label_name(str(cls_name))
next_id = max(next_id, int(cls_id) + 1)
if not id2label:
raise RuntimeError(f"Labelmap vazio ou inválido: {labelmap_path}")
ids_sorted = sorted(id2label.keys())
if ids_sorted != list(range(len(ids_sorted))):
remap = {old_id: new_id for new_id, old_id in enumerate(ids_sorted)}
id2label = {remap[old_id]: name for old_id, name in id2label.items()}
label2id = {name: idx for idx, name in id2label.items()}
return id2label, label2id, ignore_index
def get_visual_mode(config: dict) -> str:
use_mask2 = bool(config.get("dual_head_mask", config.get("dual_head", False)))
use_label = bool(config.get("dual_head_label", False))
if use_mask2 and use_label:
raise RuntimeError("Config inválido: dual_head_mask e dual_head_label ativos juntos.")
if use_label:
return "label"
if use_mask2:
return "mask2"
return "single"
def get_save_suffix(mode: str) -> str:
if mode == "single":
return "_single"
if mode == "mask2":
return "_dual_mask"
if mode == "label":
return "_dual_label"
raise RuntimeError(f"Modo desconhecido: {mode}")
def resolve_labelmap_path(args, config: dict, config_dir: Path) -> Path:
explicit = resolve_path(args.labelmap, Path.cwd())
if explicit is not None:
return explicit.resolve()
camera = str(config.get("camera", "oak-d"))
return (config_dir / camera / "dataset" / "labelmap.txt").resolve()
def resolve_default_paths(args, config: dict, config_dir: Path) -> Tuple[Path, Path, str, str, Path]:
camera = str(config.get("camera", "oak-d"))
model_key = str(config.get("modelo", "segformer_b0"))
model_name = str(config.get("model_name", "visual"))
mode = get_visual_mode(config)
suffix = get_save_suffix(mode)
save_dir = config_dir / camera / "backup" / model_key / f"{model_name}{suffix}"
if mode == "label":
default_ckpt_name = "best_label"
elif mode == "mask2":
default_ckpt_name = "best_mask2"
else:
default_ckpt_name = "best_main"
ckpt_name = str(config.get("ckpt_test", default_ckpt_name))
checkpoint_path = resolve_path(args.checkpoint, Path.cwd())
if checkpoint_path is None:
checkpoint_path = save_dir / f"{ckpt_name}.pt"
onnx_path = resolve_path(args.onnx, Path.cwd())
if onnx_path is None:
onnx_path = save_dir / f"{ckpt_name}.onnx"
if 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 not onnx_path.is_file():
raise FileNotFoundError(
f"ONNX não encontrado: {onnx_path}\n"
f"Dica: informe --onnx ou exporte antes com _10_export_visual_onnx.py."
)
return checkpoint_path.resolve(), onnx_path.resolve(), ckpt_name, mode, save_dir.resolve()
def resolve_norm_stats_path(args, config: dict, config_dir: Path, save_dir: Path) -> Optional[Path]:
explicit = resolve_path(args.norm_stats, Path.cwd())
if explicit is not None:
return explicit.resolve()
p = save_dir / "norm_stats.json"
if p.is_file():
return p.resolve()
W, H = config.get("resolucao", [1024, 640])
camera = str(config.get("camera", "oak-d"))
p = config_dir / camera / "dataset" / f"{int(W)}x{int(H)}" / "group" / "norm_stats.json"
if p.is_file():
return p.resolve()
return p.resolve()
def load_rgb_norm_stats(path: Optional[Path]) -> Tuple[Optional[List[float]], Optional[List[float]], Optional[str], List[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, ["R", "G", "B"]
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}")
if names:
name_to_idx = {str(n).upper(): i for i, n in enumerate(names)}
required = ["R", "G", "B"]
missing = [ch for ch in required if ch not in name_to_idx]
if missing:
raise RuntimeError(f"norm_stats incompatível: faltam canais {missing}. channels={names}")
idx = [name_to_idx[ch] for ch in required]
mean_sel = [float(mean[i]) for i in idx]
std_sel = [float(std[i]) for i in idx]
names_sel = required
else:
if len(mean) < 3 or len(std) < 3:
raise RuntimeError(f"norm_stats precisa de pelo menos 3 valores RGB: {path}")
mean_sel = [float(mean[i]) for i in range(3)]
std_sel = [float(std[i]) for i in range(3)]
names_sel = ["R", "G", "B"]
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), names_sel
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 resolve_label_classes(config: dict, ckpt: Optional[dict] = None) -> Tuple[Dict[int, str], int]:
label_classes = config.get("label_classes", None)
if label_classes is not None:
label_name_by_id = {i: str(name) for i, name in enumerate(label_classes)}
return label_name_by_id, len(label_name_by_id)
if ckpt is not None:
extra = ckpt.get("extra", {}) if isinstance(ckpt, dict) else {}
maybe = extra.get("label_name_by_id", None)
if isinstance(maybe, dict) and maybe:
label_name_by_id = {int(k): str(v) for k, v in maybe.items()}
return label_name_by_id, max(label_name_by_id.keys()) + 1
maybe_config = extra.get("config", {}) if isinstance(extra, dict) else {}
maybe_classes = maybe_config.get("label_classes", None) if isinstance(maybe_config, dict) else None
if maybe_classes is not None:
label_name_by_id = {i: str(name) for i, name in enumerate(maybe_classes)}
return label_name_by_id, len(label_name_by_id)
raise RuntimeError(
"Não consegui resolver label_classes. "
"Adicione config['label_classes'] ou use um checkpoint com extra['label_name_by_id']."
)
# ============================================================
# Dataset / imagens
# ============================================================
@dataclass
class Sample:
img_path: str
group_name: str
filename: str
IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp")
def discover_samples(split_root: Path, max_samples: int = 50, start_idx: int = 0) -> List[Sample]:
split_root = Path(split_root)
group_root = split_root / "group"
if not group_root.is_dir():
raise RuntimeError(f"Não achei pasta: {group_root}")
samples: List[Sample] = []
img_dirs = glob.glob(str(group_root / "**" / "images"), recursive=True)
img_dirs = [d for d in img_dirs if os.path.isdir(d)]
for idir in sorted(img_dirs):
base = os.path.dirname(idir)
group_name = os.path.relpath(base, str(group_root)).replace("\\", "/")
img_paths: List[str] = []
for ext in IMG_EXTS:
img_paths.extend(glob.glob(os.path.join(idir, f"*{ext}")))
img_paths.extend(glob.glob(os.path.join(idir, f"*{ext.upper()}")))
for ip in sorted(set(img_paths)):
samples.append(Sample(
img_path=ip,
group_name=group_name,
filename=os.path.basename(ip),
))
if not samples:
raise RuntimeError(f"Nenhuma imagem encontrada em: {group_root}/**/images")
start_idx = max(0, int(start_idx))
selected = samples[start_idx:]
if max_samples > 0:
selected = selected[:int(max_samples)]
return selected
def discover_image_folder(folder: Path, max_samples: int = 50, start_idx: int = 0) -> List[Sample]:
folder = Path(folder)
samples: List[Sample] = []
img_paths: List[str] = []
for ext in IMG_EXTS:
img_paths.extend(glob.glob(str(folder / f"*{ext}")))
img_paths.extend(glob.glob(str(folder / f"*{ext.upper()}")))
for ip in sorted(set(img_paths)):
samples.append(Sample(
img_path=ip,
group_name="external",
filename=os.path.basename(ip),
))
if not samples:
raise RuntimeError(f"Nenhuma imagem encontrada em: {folder}")
start_idx = max(0, int(start_idx))
selected = samples[start_idx:]
if max_samples > 0:
selected = selected[:int(max_samples)]
return selected
def load_rgb_image(path: str | Path) -> np.ndarray:
img_bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
if img_bgr is None:
raise RuntimeError(f"Falha ao ler imagem: {path}")
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
chw = np.transpose(img_rgb.astype(np.float32), (2, 0, 1)) / 255.0
return np.clip(chw, 0.0, 1.0).astype(np.float32)
def resize_chw(chw: np.ndarray, target_hw: Tuple[int, int]) -> np.ndarray:
H, W = target_hw
if chw.shape[-2:] == (H, W):
return chw.astype(np.float32, copy=False)
hwc = np.transpose(chw, (1, 2, 0))
hwc = cv2.resize(hwc, (W, H), interpolation=cv2.INTER_AREA)
return np.transpose(hwc, (2, 0, 1)).astype(np.float32)
def load_inputs_as_numpy(
samples: List[Sample],
mean: Optional[List[float]],
std: Optional[List[float]],
target_hw: Tuple[int, int],
normalize_input: bool = True,
) -> List[np.ndarray]:
xs: List[np.ndarray] = []
for s in samples:
chw01 = load_rgb_image(s.img_path)
chw01 = resize_chw(chw01, target_hw=target_hw)
if normalize_input:
chw = normalize_numpy_chw(chw01, mean=mean, std=std)
else:
chw = chw01.astype(np.float32, copy=False)
xs.append(np.expand_dims(chw, axis=0).astype(np.float32))
return xs
# ============================================================
# Modelo PyTorch visual
# ============================================================
class LabelHead(nn.Module):
def __init__(
self,
feat_ch: int,
num_seg_classes: int,
num_label_classes: int,
hidden: int = 256,
dropout: float = 0.2,
):
super().__init__()
in_ch = int(feat_ch) + int(num_seg_classes)
self.in_ch = in_ch
self.pool = nn.AdaptiveAvgPool2d((1, 1))
self.net = nn.Sequential(
nn.Linear(in_ch, hidden),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
nn.Linear(hidden, num_label_classes),
)
def forward(self, feat: torch.Tensor, logits_seg: torch.Tensor) -> torch.Tensor:
feat = F.interpolate(
feat,
size=logits_seg.shape[-2:],
mode="bilinear",
align_corners=False,
)
x = torch.cat([feat, logits_seg], dim=1)
x = self.pool(x).flatten(1)
return self.net(x)
def get_last_feat(out, logits: torch.Tensor) -> torch.Tensor:
if hasattr(out, "hidden_states") and out.hidden_states is not None:
feat = out.hidden_states[-1]
else:
feat = logits
feat = F.interpolate(
feat,
size=logits.shape[-2:],
mode="bilinear",
align_corners=False,
)
return feat
class VisualSegformerDualLabel(nn.Module):
def __init__(self, base_model: nn.Module, label_head: nn.Module):
super().__init__()
self.base_model = base_model
self.label_head = label_head
def forward(self, pixel_values: torch.Tensor) -> Dict[str, torch.Tensor]:
out = self.base_model(pixel_values=pixel_values)
logits_seg = out.logits
feat = get_last_feat(out, logits_seg)
logits_label = self.label_head(feat, logits_seg)
return {
"semantic_logits": logits_seg,
"label_logits": logits_label,
"label_probs": torch.softmax(logits_label, dim=1),
}
def build_visual_model(
backbone: str,
num_seg_classes: int,
num_label_classes: int,
device: torch.device,
input_hw: Tuple[int, int],
) -> VisualSegformerDualLabel:
H, W = input_hw
base_model = SegformerForSemanticSegmentation.from_pretrained(
backbone,
num_labels=int(num_seg_classes),
ignore_mismatched_sizes=True,
use_safetensors=True,
)
base_model.config.output_hidden_states = True
base_model.to(device)
base_model.eval()
with torch.no_grad():
dummy = torch.zeros((1, 3, int(H), int(W)), dtype=torch.float32, device=device)
out = base_model(pixel_values=dummy)
logits = out.logits
feat = get_last_feat(out, logits)
feat_ch = int(feat.shape[1])
label_head = LabelHead(
feat_ch=feat_ch,
num_seg_classes=int(num_seg_classes),
num_label_classes=int(num_label_classes),
hidden=256,
dropout=0.2,
).to(device)
label_head.eval()
return VisualSegformerDualLabel(base_model=base_model, label_head=label_head).to(device)
def load_checkpoint_into_model(model: VisualSegformerDualLabel, checkpoint_path: Path):
ckpt = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
if "model" not in ckpt:
raise RuntimeError(f"Checkpoint não contém chave 'model': {checkpoint_path}")
if "aux_head" not in ckpt:
raise RuntimeError(f"Checkpoint não contém chave 'aux_head': {checkpoint_path}")
model.base_model.load_state_dict(ckpt["model"], strict=True)
model.label_head.load_state_dict(ckpt["aux_head"], strict=True)
return ckpt
class VisualTorchTupleWrapper(nn.Module):
def __init__(self, model: VisualSegformerDualLabel, output_kind: str = "contract"):
super().__init__()
self.model = model
self.output_kind = str(output_kind).lower()
if self.output_kind not in ("raw", "contract"):
raise RuntimeError(f"output_kind inválido: {self.output_kind}")
def forward(self, pixel_values: torch.Tensor):
outputs = self.model(pixel_values=pixel_values)
semantic = outputs["semantic_logits"]
label_probs = outputs["label_probs"]
if self.output_kind == "contract":
semantic = F.interpolate(
semantic,
size=pixel_values.shape[-2:],
mode="bilinear",
align_corners=False,
)
return semantic, label_probs
# ============================================================
# Benchmark PyTorch
# ============================================================
@torch.inference_mode()
def benchmark_torch(
model: nn.Module,
inputs_np: List[np.ndarray],
device: torch.device,
warmup: int,
repeat: int,
amp: bool,
label: str,
) -> Tuple[dict, List[dict]]:
model.eval()
times: List[float] = []
rows: List[dict] = []
inputs_t = [
torch.from_numpy(x).to(device, non_blocking=True)
for x in inputs_np
]
synchronize_if_cuda(device)
print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")
for i in range(max(0, warmup)):
x = inputs_t[i % len(inputs_t)]
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
_ = model(x)
synchronize_if_cuda(device)
total_iter = len(inputs_t) * max(1, repeat)
idx = 0
for r in range(max(1, repeat)):
for sample_idx, x in enumerate(inputs_t):
synchronize_if_cuda(device)
t0 = time.perf_counter()
with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=amp and device.type == "cuda"):
_ = model(x)
synchronize_if_cuda(device)
dt_ms = (time.perf_counter() - t0) * 1000.0
times.append(dt_ms)
rows.append({
"engine": label,
"repeat": r,
"sample_idx": sample_idx,
"iter_idx": idx,
"latency_ms": dt_ms,
})
idx += 1
if idx % 25 == 0 or idx == total_iter:
print(f" {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")
summary = summarize_times(times)
summary["engine"] = label
return summary, rows
# ============================================================
# ONNX Runtime
# ============================================================
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 CPU:\n"
" pip install onnxruntime"
)
provider = provider.lower()
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"))
dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\bin")
dll_dirs.append(r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.3\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),
"trt_engine_cache_enable": True,
"trt_engine_cache_path": str(cache_dir),
"trt_timing_cache_enable": True,
"trt_timing_cache_path": str(cache_dir),
"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}"
)
session = ort.InferenceSession(
str(onnx_path),
sess_options=sess_options,
providers=providers_ok,
)
active = session.get_providers()
print(f"[ONNX] usando providers: {active}")
if provider == "tensorrt" and "TensorrtExecutionProvider" not in active:
raise RuntimeError(
"TensorRTExecutionProvider foi solicitado, mas não ficou ativo. "
f"Providers ativos: {active}."
)
if provider == "cuda" and "CUDAExecutionProvider" not in active:
raise RuntimeError(
"CUDAExecutionProvider foi solicitado, mas não ficou ativo. "
f"Providers ativos: {active}."
)
return session
def benchmark_onnx(
session,
inputs_np: List[np.ndarray],
warmup: int,
repeat: int,
label: str,
) -> Tuple[dict, List[dict]]:
input_name = session.get_inputs()[0].name
times: List[float] = []
rows: List[dict] = []
print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}")
for i in range(max(0, warmup)):
x = inputs_np[i % len(inputs_np)]
_ = session.run(None, {input_name: x})
total_iter = len(inputs_np) * max(1, repeat)
idx = 0
for r in range(max(1, repeat)):
for sample_idx, x in enumerate(inputs_np):
t0 = time.perf_counter()
_ = session.run(None, {input_name: x})
dt_ms = (time.perf_counter() - t0) * 1000.0
times.append(dt_ms)
rows.append({
"engine": label,
"repeat": r,
"sample_idx": sample_idx,
"iter_idx": idx,
"latency_ms": dt_ms,
})
idx += 1
if idx % 25 == 0 or idx == total_iter:
print(f" {idx:04d}/{total_iter:04d} | last={dt_ms:.2f}ms")
summary = summarize_times(times)
summary["engine"] = label
return summary, rows
# ============================================================
# 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("--labelmap", default="")
parser.add_argument("--norm_stats", default="")
parser.add_argument("--split_folder", default="val", choices=["train", "val", "test"])
parser.add_argument("--root_override", default=None)
parser.add_argument("--test_folder", default=None)
parser.add_argument("--max_samples", type=int, default=50)
parser.add_argument("--start_idx", type=int, default=0)
parser.add_argument("--warmup", type=int, default=10)
parser.add_argument("--repeat", type=int, default=5)
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")
parser.add_argument("--trt_no_fp16", action="store_true")
parser.add_argument("--skip_torch_fp32", action="store_true")
parser.add_argument("--skip_torch_amp", action="store_true")
parser.add_argument("--skip_onnx", action="store_true")
parser.add_argument(
"--onnx_has_norm",
action="store_true",
help="Use quando o ONNX já inclui normalização interna. Nesse caso o ONNX recebe RGB 0..1 cru.",
)
parser.add_argument(
"--torch_output_kind",
default="contract",
choices=["raw", "contract"],
help="contract mede PyTorch com semantic_logits redimensionado para HxW, igual ao ONNX exportado com resize_logits.",
)
parser.add_argument("--out_dir", default=None)
args = parser.parse_args()
config_path = resolve_path(args.config, Path.cwd())
if config_path is None or not config_path.is_file():
raise FileNotFoundError(f"Config não encontrado: {config_path}")
config_dir = config_path.parent
config = load_json(config_path)
checkpoint_path, onnx_path, ckpt_name, mode, save_dir = resolve_default_paths(args, config, config_dir)
if mode != "label":
raise RuntimeError(f"Este benchmark foi preparado para dual_head_label. Modo detectado: {mode}")
labelmap_path = resolve_labelmap_path(args, config, config_dir)
semantic_id2label, semantic_label2id, ignore_index = load_labelmap(labelmap_path)
num_seg_classes = len(semantic_id2label)
W, H = config.get("resolucao", [1024, 640])
W = int(W)
H = int(H)
backbone = str(config.get("backbone", "nvidia/mit-b0"))
ckpt_meta = torch.load(str(checkpoint_path), map_location="cpu", weights_only=False)
label_name_by_id, num_label_classes = resolve_label_classes(config, ckpt_meta)
norm_stats_path = resolve_norm_stats_path(args, config, config_dir, save_dir)
mean, std, norm_stats_used, norm_channels = load_rgb_norm_stats(norm_stats_path)
if args.test_folder:
root = resolve_path(args.test_folder, Path.cwd())
if root is None or not root.is_dir():
raise FileNotFoundError(f"Pasta de teste não encontrada: {root}")
samples = discover_image_folder(
folder=root,
max_samples=args.max_samples,
start_idx=args.start_idx,
)
else:
if args.root_override:
root = resolve_path(args.root_override, Path.cwd())
else:
camera = str(config.get("camera", "oak-d"))
root = (config_dir / camera / "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 = discover_samples(
split_root=root,
max_samples=args.max_samples,
start_idx=args.start_idx,
)
if args.out_dir:
out_dir = resolve_path(args.out_dir, Path.cwd())
else:
out_dir = onnx_path.parent / "benchmarks"
assert out_dir is not None
out_dir.mkdir(parents=True, exist_ok=True)
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("Benchmark Visual Worker 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"Warmup : {args.warmup}")
print(f"Repeat : {args.repeat}")
print(f"Backbone : {backbone}")
print(f"Input shape : [1, 3, {H}, {W}]")
print(f"Semantic classes: {num_seg_classes} {semantic_id2label}")
print(f"Label classes : {num_label_classes} {label_name_by_id}")
print(f"Device : {device}")
print(f"ONNX provider : {args.onnx_provider}")
print(f"ONNX has norm : {args.onnx_has_norm}")
print(f"Torch output : {args.torch_output_kind}")
print(f"Out dir : {out_dir}")
print("==========================================")
# Entradas PyTorch: sempre normalizadas, porque o modelo PyTorch puro espera normalização externa.
# Entradas ONNX: se ONNX tem norm embutida, entram 0..1; senão, entram normalizadas.
print("\n[DATA] Carregando imagens na RAM...")
inputs_torch_np = load_inputs_as_numpy(
samples=samples,
mean=mean,
std=std,
target_hw=(H, W),
normalize_input=True,
)
if args.onnx_has_norm:
inputs_onnx_np = load_inputs_as_numpy(
samples=samples,
mean=mean,
std=std,
target_hw=(H, W),
normalize_input=False,
)
print("[DATA] ONNX receberá RGB 0..1 cru, pois --onnx_has_norm está ativo.")
else:
inputs_onnx_np = inputs_torch_np
print("[DATA] ONNX receberá RGB normalizado, pois --onnx_has_norm não está ativo.")
print(f"[DATA] Inputs carregados: {len(inputs_torch_np)}")
summaries: List[dict] = []
all_rows: List[dict] = []
# ========================================================
# PyTorch
# ========================================================
need_torch = not args.skip_torch_fp32 or not args.skip_torch_amp
if need_torch:
print("\n[MODEL] Montando PyTorch...")
model = build_visual_model(
backbone=backbone,
num_seg_classes=num_seg_classes,
num_label_classes=num_label_classes,
device=device,
input_hw=(H, W),
)
print("[CKPT] Carregando checkpoint...")
ckpt = load_checkpoint_into_model(model, checkpoint_path)
model.to(device)
model.eval()
torch_model = VisualTorchTupleWrapper(
model=model,
output_kind=args.torch_output_kind,
).to(device)
torch_model.eval()
clear_cuda()
if not args.skip_torch_fp32:
summary, rows = benchmark_torch(
model=torch_model,
inputs_np=inputs_torch_np,
device=device,
warmup=args.warmup,
repeat=args.repeat,
amp=False,
label="torch_fp32",
)
summaries.append(summary)
all_rows.extend(rows)
clear_cuda()
if not args.skip_torch_amp:
summary, rows = benchmark_torch(
model=torch_model,
inputs_np=inputs_torch_np,
device=device,
warmup=args.warmup,
repeat=args.repeat,
amp=True,
label="torch_amp_fp16",
)
summaries.append(summary)
all_rows.extend(rows)
del torch_model
del model
clear_cuda()
# ========================================================
# ONNX
# ========================================================
if not args.skip_onnx:
print("\n[ONNX] Carregando sessão...")
session = create_onnx_session(
onnx_path=onnx_path,
provider=args.onnx_provider,
trt_home=args.trt_home,
trt_fp16=not args.trt_no_fp16,
)
summary, rows = benchmark_onnx(
session=session,
inputs_np=inputs_onnx_np,
warmup=args.warmup,
repeat=args.repeat,
label=f"onnx_{args.onnx_provider}",
)
summaries.append(summary)
all_rows.extend(rows)
# ========================================================
# Relatório
# ========================================================
print("\n========== RESUMO ==========")
for s in summaries:
print(
f"{s['engine']:<16} "
f"n={s['n']:<5} "
f"mean={s['mean_ms']:.3f}ms "
f"median={s['median_ms']:.3f}ms "
f"p95={s['p95_ms']:.3f}ms "
f"p99={s['p99_ms']:.3f}ms "
f"fps_mean={s['fps_mean']:.2f} "
f"fps_p95={s['fps_p95_latency']:.2f}"
)
base_name = f"{onnx_path.stem}_{args.onnx_provider}"
report_json = out_dir / f"{base_name}_visual_benchmark_report.json"
report_csv = out_dir / f"{base_name}_visual_benchmark_rows.csv"
report = {
"kind": "visual_worker_benchmark_onnx",
"config": str(config_path),
"checkpoint": str(checkpoint_path),
"ckpt_name": ckpt_name,
"onnx": str(onnx_path),
"onnx_has_norm": bool(args.onnx_has_norm),
"root": str(root),
"samples": len(samples),
"warmup": int(args.warmup),
"repeat": int(args.repeat),
"input_shape": [1, 3, H, W],
"input_channel_names": ["R", "G", "B"],
"semantic_id2label": semantic_id2label,
"label_name_by_id": label_name_by_id,
"norm_stats_used": norm_stats_used,
"norm_channels": norm_channels,
"onnx_provider": args.onnx_provider,
"trt_fp16": bool(not args.trt_no_fp16),
"device": str(device),
"torch_output_kind": args.torch_output_kind,
"summaries": summaries,
"samples_list": [
{
"img_path": s.img_path,
"group_name": s.group_name,
"filename": s.filename,
}
for s in samples
],
}
save_json(report_json, report)
save_csv(report_csv, all_rows)
print(f"\n[OK] JSON salvo em: {report_json}")
print(f"[OK] CSV salvo em : {report_csv}")
print("\nBenchmark finalizado.")
if __name__ == "__main__":
main()