#!/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()