1171 lines
38 KiB
Python
1171 lines
38 KiB
Python
#!/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()
|