#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ _12_benchmark_onnx.py Benchmark PyTorch vs ONNX Runtime para SegFormer OAK-FCC-3 Multi-Head. Mede: - PyTorch FP32 - PyTorch AMP/FP16 - ONNX Runtime CUDA ou CPU Exemplos: Benchmark ONNX cru 160x256: python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda Benchmark ONNX resized 640x1024: python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider cuda TensorRT python _12_benchmark_onnx.py --config config.json --max_samples 50 --warmup 10 --repeat 5 --device cuda --onnx_provider tensorrt --skip_torch_fp32 --skip_torch_amp """ from __future__ import annotations import gc import csv import json import time import argparse import importlib.util from pathlib import Path from typing import Optional, List, Dict, Sequence, Tuple import cv2 import numpy as np import torch import torch.nn as nn import torch.nn.functional as F SOURCE_CHANNEL_ORDER = ["R", "G", "B", "RE", "NIR"] DERIVED_CHANNEL_ORDER = ["NDVI", "NDRE"] SUPPORTED_INPUT_CHANNELS = SOURCE_CHANNEL_ORDER + DERIVED_CHANNEL_ORDER # ============================================================ # 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 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: 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 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 percentile(values: List[float], p: float) -> float: if not values: return 0.0 return float(np.percentile(np.asarray(values, dtype=np.float64), p)) 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, "p95_ms": 0.0, "p99_ms": 0.0, "fps_mean": 0.0, "fps_p95_latency": 0.0, } mean_ms = float(arr.mean()) 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": float(np.median(arr)), "min_ms": float(arr.min()), "max_ms": float(arr.max()), "p95_ms": p95_ms, "p99_ms": p99_ms, "fps_mean": float(1000.0 / mean_ms) if mean_ms > 0 else 0.0, "fps_p95_latency": float(1000.0 / p95_ms) if p95_ms > 0 else 0.0, } def experiment_tag(config: dict, channels: int) -> str: explicit = str(config.get("stats_source_tag", "")).strip() if explicit: return explicit return f"{config.get('fusion_mode', 'stacked')}_raw{channels}" def load_export_metadata(onnx_path: Path) -> Tuple[dict, Optional[Path]]: meta_path = onnx_path.with_suffix(".export_meta.json") if meta_path.is_file(): return load_json(meta_path), meta_path return {}, None 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}/{stats_source_tag}/{ckpt_name}.pt backup/{modelo}/{model_name}/{stats_source_tag}/{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") ckpt_name = config.get("ckpt_test", "best_score") base_dir = config_dir / "backup" / model / model_name / experiment_tag(config, channels) ckpt_filename = str(ckpt_name) if not ckpt_filename.lower().endswith(".pt"): ckpt_filename += ".pt" onnx_filename = f"{Path(str(ckpt_name)).stem}.onnx" if args.checkpoint: checkpoint_path = resolve_path(args.checkpoint, config_dir) else: checkpoint_path = base_dir / ckpt_filename if args.onnx: onnx_path = resolve_path(args.onnx, config_dir) else: onnx_path = base_dir / onnx_filename 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) # ============================================================ # Dataset / normalização # ============================================================ def collect_tensor_samples(root: Path, max_samples: int = 50, 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 load_tensor( path: Path, requested_channels: Sequence[str], saved_channel_names: Optional[Sequence[str]] = None, ) -> 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 1 <= arr.shape[0] <= len(SUPPORTED_INPUT_CHANNELS) and arr.shape[1] > 8 and arr.shape[2] > 8: chw = arr elif 1 <= arr.shape[-1] <= len(SUPPORTED_INPUT_CHANNELS) and arr.shape[0] > 8 and arr.shape[1] > 8: chw = np.transpose(arr, (2, 0, 1)) else: raise RuntimeError(f"Tensor com layout inesperado: {path} shape={arr.shape}") requested = [str(x).strip().upper() for x in requested_channels] saved = [str(x).strip().upper() for x in (saved_channel_names or [])] if saved: if len(saved) != chw.shape[0] or len(set(saved)) != len(saved): raise RuntimeError( f"Meta/tensor incompatíveis em {path}: channels={saved}, shape={chw.shape}" ) missing = [name for name in requested if name not in saved] if missing: raise RuntimeError(f"Tensor {path} não contém {missing}. Disponíveis={saved}") indices = [saved.index(name) for name in requested] elif chw.shape[0] == len(requested): indices = list(range(len(requested))) print(f"[WARN] {path.name}: sem meta.channels; assumindo ordem {requested}") elif chw.shape[0] == len(SOURCE_CHANNEL_ORDER) and all(x in SOURCE_CHANNEL_ORDER for x in requested): indices = [SOURCE_CHANNEL_ORDER.index(name) for name in requested] print(f"[WARN] {path.name}: Raw5 legado; usando ordem {SOURCE_CHANNEL_ORDER}") else: raise RuntimeError( f"Não é seguro inferir canais de {path}: shape={chw.shape}, pedidos={requested}" ) chw = chw[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=-1.0) return np.ascontiguousarray(chw, dtype=np.float32) def load_tensor_channel_names(tensor_path: Path) -> List[str]: meta_path = tensor_path.parent.parent / "metas" / f"{tensor_path.stem}.json" if not meta_path.is_file(): return [] meta = load_json(meta_path) return list(meta.get("channels") or meta.get("input_channels") or []) def find_norm_stats(config: dict, config_dir: Path, save_dir: Path, explicit: Optional[str]) -> Optional[Path]: if explicit: return resolve_path(explicit, config_dir) W, H = config.get("resolucao", [1024, 640]) dataset_path = config_dir / "dataset" candidates = [ save_dir / "norm_stats.json", dataset_path / f"{int(W)}x{int(H)}" / "group" / "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 load_norm_stats( path: Optional[Path], channel_names: Sequence[str], ) -> Tuple[Optional[List[float]], Optional[List[float]], Optional[str]]: if path is None or not path.is_file(): raise FileNotFoundError( f"norm_stats não encontrado: {path}. Benchmark exige os stats do treinamento." ) js = load_json(path) mean = js.get("mean") std = js.get("std") names = [str(x).strip().upper() for x in js.get("channels", [])] requested = [str(x).strip().upper() for x in channel_names] if mean is None or std is None: raise RuntimeError(f"norm_stats inválido, faltando mean/std: {path}") if len(mean) != len(std): raise RuntimeError(f"norm_stats inválido: mean={len(mean)} std={len(std)}") if names: if len(names) != len(mean) or len(set(names)) != len(names): raise RuntimeError(f"norm_stats channels inválido: {names}") missing = [name for name in requested if name not in names] if missing: raise RuntimeError( f"norm_stats não contém {missing}. Disponíveis={names}, pedidos={requested}" ) indices = [names.index(name) for name in requested] elif len(mean) == len(requested): indices = list(range(len(requested))) print(f"[NORM][WARN] stats sem nomes; assumindo ordem {requested}") elif len(mean) == len(SOURCE_CHANNEL_ORDER) and all(x in SOURCE_CHANNEL_ORDER for x in requested): indices = [SOURCE_CHANNEL_ORDER.index(name) for name in requested] print(f"[NORM][WARN] stats Raw5 legados; usando ordem {SOURCE_CHANNEL_ORDER}") else: raise RuntimeError( f"Não é seguro mapear norm_stats sem nomes: mean={len(mean)}, pedidos={requested}" ) mean_sel = [float(mean[i]) for i in indices] std_sel = [float(std[i]) for i in indices] if not np.all(np.isfinite(mean_sel)) or not np.all(np.isfinite(std_sel)): raise RuntimeError(f"norm_stats possui NaN/Inf: mean={mean_sel} std={std_sel}") if any(x <= 0.0 for x in std_sel): raise RuntimeError(f"norm_stats possui std inválido: {std_sel}") print(f"[NORM] usando {path}") print(f"[NORM] channels={requested}") 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_inputs_as_numpy( samples: List[Path], input_channel_names: Sequence[str], mean: Optional[List[float]], std: Optional[List[float]], target_hw: Tuple[int, int], normalize_input: bool = True, ) -> List[np.ndarray]: H, W = target_hw xs = [] for p in samples: chw01 = load_tensor( p, requested_channels=input_channel_names, saved_channel_names=load_tensor_channel_names(p), ) 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) if normalize_input: chw = normalize_numpy_chw(chw01, mean=mean, std=std) else: chw = chw01.astype(np.float32, copy=False) x = np.expand_dims(chw, axis=0).astype(np.float32) xs.append(x) return xs # ============================================================ # PyTorch # ============================================================ class TorchTupleWrapper(nn.Module): def __init__( self, model: nn.Module, output_heads: List[str], postprocess: str = "none", include_norm: bool = False, norm_mean: Optional[Sequence[float]] = None, norm_std: Optional[Sequence[float]] = None, ): super().__init__() self.model = model self.output_heads = list(output_heads) self.postprocess = str(postprocess).lower() self.include_norm = bool(include_norm) if self.postprocess not in ("none", "resize_logits", "argmax_lowres", "argmax_fullres"): raise RuntimeError(f"postprocess inválido: {self.postprocess}") if self.include_norm: if norm_mean is None or norm_std is None: raise RuntimeError("include_norm=True requer norm_mean/norm_std") self.register_buffer( "norm_mean", torch.tensor(norm_mean, dtype=torch.float32).view(1, -1, 1, 1), ) self.register_buffer( "norm_std", torch.tensor(norm_std, dtype=torch.float32).view(1, -1, 1, 1), ) else: self.register_buffer("norm_mean", torch.empty(0)) self.register_buffer("norm_std", torch.empty(0)) def forward(self, pixel_values: torch.Tensor): x = pixel_values if self.include_norm: x = (x - self.norm_mean) / torch.clamp(self.norm_std, min=1e-6) outputs = self.model(pixel_values=x) input_hw = pixel_values.shape[-2:] result = [] for head in self.output_heads: value = outputs[head] if self.postprocess in ("resize_logits", "argmax_fullres"): value = F.interpolate( value, size=input_hw, mode="bilinear", align_corners=False, ) if self.postprocess in ("argmax_lowres", "argmax_fullres"): value = torch.argmax(value, dim=1).to(torch.uint8) result.append(value) return tuple(result) @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 = [] rows = [] # Precarrega tensors na GPU para medir só inferência do modelo. inputs_t = [ torch.from_numpy(x).to(device, non_blocking=True) for x in inputs_np ] if device.type == "cuda": torch.cuda.synchronize() print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}") # Warmup 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) # Medição 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 summary["timing_scope"] = "model_forward_and_export_postprocess_device_resident" return summary, rows # ============================================================ # ONNX Runtime # ============================================================ def create_onnx_session( onnx_path: Path, provider: str, trt_home: Optional[str] = None, trt_fp16: bool = True, ): import os 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" ) available = ort.get_available_providers() print(f"[ONNX] providers disponíveis: {available}") 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") 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}") 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, # FP16: o ponto principal do nosso teste. "trt_fp16_enable": bool(trt_fp16), # Cache: evita rebuild do engine a cada execução. "trt_engine_cache_enable": True, "trt_engine_cache_path": str(cache_dir), # Timing cache ajuda a acelerar builds futuros. "trt_timing_cache_enable": True, "trt_timing_cache_path": str(cache_dir), # Workspace. 4GB é razoável para RTX 3070, ajuste se faltar VRAM. "trt_max_workspace_size": 4 * 1024 * 1024 * 1024, } providers = [ ("TensorrtExecutionProvider", trt_options), "CUDAExecutionProvider", "CPUExecutionProvider", ] else: providers = [provider] # Checagem de disponibilidade, lidando com provider tuple. 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, ) print(f"[ONNX] usando providers: {session.get_providers()}") active_providers = session.get_providers() if provider == "tensorrt" and "TensorrtExecutionProvider" not in active_providers: raise RuntimeError( "TensorRTExecutionProvider foi solicitado, mas não ficou ativo. " f"Providers ativos: {active_providers}. " "Provável causa: TensorRT não instalado, DLLs fora do PATH, " "ou versão incompatível com onnxruntime-gpu." ) if provider == "cuda" and "CUDAExecutionProvider" not in active_providers: raise RuntimeError( "CUDAExecutionProvider foi solicitado, mas não ficou ativo. " f"Providers ativos: {active_providers}." ) 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 = [] rows = [] print(f"\n[BENCH] {label} | warmup={warmup} repeat={repeat}") # Warmup for i in range(max(0, warmup)): x = inputs_np[i % len(inputs_np)] _ = session.run(None, {input_name: x}) # Medição 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 summary["timing_scope"] = "onnx_session_run_host_input_and_outputs" 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("--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=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=None) 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="Força ONNX com normalização interna; normalmente detectado pelo export_meta.", ) parser.add_argument( "--onnx_no_norm", action="store_true", help="Força ONNX sem normalização interna; normalmente detectado pelo export_meta.", ) parser.add_argument("--out_dir", default=None) args = parser.parse_args() if args.skip_torch_fp32 and args.skip_torch_amp and args.skip_onnx: raise RuntimeError("Todos os engines foram desativados; não há o que medir.") 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) train_script_path = resolve_path(args.train_script, config_dir) labelmap_path = resolve_path(args.labelmap, config_dir) 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}") 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, ) export_meta, export_meta_path = load_export_metadata(onnx_path) exported_names = [ str(x).strip().upper() for x in export_meta.get("input_channel_names", []) ] if exported_names and exported_names != input_channel_names: raise RuntimeError( f"ONNX incompatível: canais exportados={exported_names}, " f"config atual={input_channel_names}." ) if args.onnx_has_norm and args.onnx_no_norm: raise RuntimeError("Use apenas uma opção: --onnx_has_norm ou --onnx_no_norm.") input_contract = export_meta.get("input_contract", {}) or {} detected_onnx_has_norm = bool( input_contract.get("normalization_embedded", export_meta.get("include_norm", False)) ) onnx_has_norm = ( True if args.onnx_has_norm else False if args.onnx_no_norm else detected_onnx_has_norm ) exported_postprocess = str(export_meta.get("postprocess", "none")).lower() if exported_postprocess == "none" and bool(export_meta.get("resize_to_input", False)): exported_postprocess = "resize_logits" if exported_postprocess not in ("none", "resize_logits", "argmax_lowres", "argmax_fullres"): raise RuntimeError(f"postprocess inválido no export_meta: {exported_postprocess}") 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()) exported_heads = list(export_meta.get("heads", []) or []) if exported_heads and exported_heads != output_heads: raise RuntimeError( f"Heads do ONNX={exported_heads} diferem do config/modelo={output_heads}." ) save_dir = ( config_dir / "backup" / config.get("modelo", "segformer_b1") / config.get("model_name", "test") / experiment_tag(config, 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, 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, ) if args.out_dir: out_dir = resolve_path(args.out_dir, config_dir) 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 PyTorch vs ONNX") print(f"Config : {config_path}") print(f"Checkpoint : {checkpoint_path}") print(f"ONNX : {onnx_path}") print(f"Export meta : {export_meta_path or 'ausente (modo legado)'}") 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, {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"ONNX has norm: {onnx_has_norm}") print(f"Postprocess : {exported_postprocess}") print(f"Out dir : {out_dir}") print("==========================================") print( "[TIMING] PyTorch mede forward com entrada residente no device; " "ONNX session.run inclui cópia da entrada/saídas do ORT. " "Pré-processamento e leitura de disco ficam fora das duas medições." ) run_torch_fp32 = not args.skip_torch_fp32 run_torch_amp = not args.skip_torch_amp and device.type == "cuda" if not args.skip_torch_amp and device.type != "cuda": print("[WARN] PyTorch AMP/FP16 exige CUDA; medição AMP será pulada.") need_torch = run_torch_fp32 or run_torch_amp if not need_torch and args.skip_onnx: raise RuntimeError("Nenhum engine executável restou para o benchmark.") need_normalized_inputs = ( (need_torch or not args.skip_onnx) and not onnx_has_norm ) inputs_torch_np: List[np.ndarray] = [] if need_normalized_inputs: print("\n[DATA] Carregando inputs normalizados...") inputs_torch_np = load_inputs_as_numpy( samples=samples, input_channel_names=input_channel_names, mean=mean, std=std, target_hw=(H, W), normalize_input=True, ) if (need_torch or not args.skip_onnx) and onnx_has_norm: print("[DATA] Carregando inputs crus para ONNX com normalização embutida...") inputs_onnx_np = load_inputs_as_numpy( samples=samples, input_channel_names=input_channel_names, mean=mean, std=std, target_hw=(H, W), normalize_input=False, ) else: inputs_onnx_np = inputs_torch_np if need_torch and onnx_has_norm: inputs_torch_np = inputs_onnx_np print(f"[DATA] Amostras preparadas: {len(samples)}") summaries = [] all_rows = [] # ======================================================== # PyTorch # ======================================================== if need_torch: print("\n[MODEL] Montando PyTorch...") model = train_mod.build_model( backbone=backbone, input_channel_names=input_channel_names, heads_config=heads_config, semantic_id2label=semantic_id2label, semantic_label2id=semantic_label2id, ) ckpt = train_mod.load_checkpoint( str(checkpoint_path), model, map_location="cpu", ) saved_contract = ckpt.get("model_contract", {}) or {} saved_names = [str(x).strip().upper() for x in saved_contract.get("input_channel_names", [])] legacy_names = SOURCE_CHANNEL_ORDER[:channels] if not saved_names and input_channel_names != legacy_names: raise RuntimeError( "Checkpoint sem contrato nominal não pode ser usado com " f"input_channels={input_channel_names}." ) model.to(device) model.eval() torch_model = TorchTupleWrapper( model=model, output_heads=output_heads, postprocess=exported_postprocess, include_norm=onnx_has_norm, norm_mean=mean, norm_std=std, ).to(device) torch_model.eval() clear_cuda() if run_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 run_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, ) onnx_input = session.get_inputs()[0] onnx_input_shape = onnx_input.shape if len(onnx_input_shape) != 4: raise RuntimeError( f"Entrada ONNX inesperada: {onnx_input.name} shape={onnx_input_shape}" ) onnx_channels = onnx_input_shape[1] if isinstance(onnx_channels, int) and onnx_channels != channels: raise RuntimeError( f"ONNX espera {onnx_channels} canais, config possui {channels}: " f"{input_channel_names}" ) print(f"[ONNX] input={onnx_input.name} shape={onnx_input_shape}") 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']:<4} " 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}_benchmark_report.json" report_csv = out_dir / f"{base_name}_benchmark_rows.csv" report = { "config": str(config_path), "checkpoint": str(checkpoint_path), "ckpt_name": ckpt_name, "onnx": str(onnx_path), "export_meta": str(export_meta_path) if export_meta_path is not None else None, "onnx_has_norm": bool(onnx_has_norm), "onnx_postprocess": exported_postprocess, "root": str(root), "samples": len(samples), "warmup": int(args.warmup), "repeat": int(args.repeat), "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, "onnx_provider": args.onnx_provider, "trt_home": args.trt_home, "trt_fp16": bool(not args.trt_no_fp16), "device": str(device), "dataset_preprocessing_in_timing": False, "normalization_in_timing": bool(onnx_has_norm), "summaries": summaries, } 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()