This commit is contained in:
ilgeco
@@ -0,0 +1,295 @@
|
||||
#!/usr/bin/env python3.13
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
VALIDATION_DIR = SCRIPT_DIR.parent
|
||||
REPO_ROOT = VALIDATION_DIR.parent
|
||||
if str(VALIDATION_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(VALIDATION_DIR))
|
||||
|
||||
from onnx_utils import _ONNX_TO_NP, onnx_io, write_inputs_to_memory_bin
|
||||
from validate_one import (
|
||||
MODE_COMPILE_ONLY,
|
||||
build_dump_ranges,
|
||||
parse_pim_simulator_outputs,
|
||||
run_pim_simulator,
|
||||
sanitize_output_name,
|
||||
validate_network,
|
||||
)
|
||||
from yolo_real_image_validation import save_tensor_csv
|
||||
|
||||
IMAGENET_MEAN = np.asarray([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
IMAGENET_STD = np.asarray([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
DEFAULT_VGG_MODEL = VALIDATION_DIR / "networks" / "vgg16" / "depth_35" / "vgg16_depth_35.onnx"
|
||||
DEFAULT_RESNET_MODEL = VALIDATION_DIR / "networks" / "resnet" / "resnet18_torchvision.onnx"
|
||||
|
||||
|
||||
def resolve_default_paths():
|
||||
return {
|
||||
"raptor_path": REPO_ROOT / "build_release" / "Release" / "bin" / "onnx-mlir",
|
||||
"onnx_include_dir": REPO_ROOT / "onnx-mlir" / "include",
|
||||
"simulator_dir": REPO_ROOT / "backend-simulators" / "pim" / "pim-simulator",
|
||||
}
|
||||
|
||||
|
||||
def resolve_model_path(network: str | None, model: Path | None) -> Path:
|
||||
if model is not None:
|
||||
return model.resolve()
|
||||
if network == "resnet":
|
||||
return DEFAULT_RESNET_MODEL.resolve()
|
||||
if network == "vgg":
|
||||
return DEFAULT_VGG_MODEL.resolve()
|
||||
raise SystemExit("Pass --model or select a default with --network {resnet,vgg}.")
|
||||
|
||||
|
||||
def ensure_local_artifacts(args, model_path: Path):
|
||||
validate_network(
|
||||
network_onnx_path=model_path,
|
||||
raptor_path=args.raptor_path,
|
||||
onnx_include_dir=args.onnx_include_dir,
|
||||
simulator_dir=args.simulator_dir,
|
||||
crossbar_size=args.crossbar_size,
|
||||
crossbar_count=args.crossbar_count,
|
||||
core_count=args.core_count,
|
||||
command_timeout_seconds=args.command_timeout_seconds,
|
||||
mode=MODE_COMPILE_ONLY,
|
||||
verbose=args.verbose,
|
||||
)
|
||||
|
||||
|
||||
def ensure_existing_artifacts(model_dir: Path):
|
||||
required_paths = [
|
||||
model_dir / "runner" / "build" / "runner",
|
||||
model_dir / "raptor" / "pim" / "config.json",
|
||||
model_dir / "raptor" / "pim" / "memory.bin",
|
||||
]
|
||||
missing = [str(path) for path in required_paths if not path.exists()]
|
||||
if missing:
|
||||
raise FileNotFoundError(
|
||||
"Missing compiled local artifacts. Re-run without --skip-compile or restore these paths:\n "
|
||||
+ "\n ".join(missing)
|
||||
)
|
||||
|
||||
|
||||
def preprocess_classification_image(image_path: Path) -> tuple[Image.Image, np.ndarray]:
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
width, height = image.size
|
||||
scale = 256.0 / min(width, height)
|
||||
resized_size = (
|
||||
max(1, int(round(width * scale))),
|
||||
max(1, int(round(height * scale))),
|
||||
)
|
||||
resized = image.resize(resized_size, Image.Resampling.BILINEAR)
|
||||
|
||||
left = (resized.width - 224) // 2
|
||||
top = (resized.height - 224) // 2
|
||||
cropped = resized.crop((left, top, left + 224, top + 224))
|
||||
|
||||
array = np.asarray(cropped, dtype=np.float32) / 255.0
|
||||
array = (array - IMAGENET_MEAN) / IMAGENET_STD
|
||||
chw = np.transpose(array, (2, 0, 1))
|
||||
tensor = np.expand_dims(chw.astype(np.float32, copy=False), axis=0)
|
||||
return image, tensor
|
||||
|
||||
|
||||
def load_labels(labels_path: Path | None) -> list[str] | None:
|
||||
if labels_path is None:
|
||||
return None
|
||||
labels = [line.strip() for line in labels_path.read_text().splitlines()]
|
||||
return labels or None
|
||||
|
||||
|
||||
def softmax(values: np.ndarray) -> np.ndarray:
|
||||
shifted = values - np.max(values)
|
||||
exp = np.exp(shifted)
|
||||
denom = exp.sum()
|
||||
if not math.isfinite(float(denom)) or denom <= 0.0:
|
||||
raise RuntimeError("Softmax received non-finite output scores.")
|
||||
return exp / denom
|
||||
|
||||
|
||||
def decode_classification_output(output: np.ndarray, labels: list[str] | None, top_k: int):
|
||||
scores = np.asarray(output, dtype=np.float64).reshape(-1)
|
||||
probabilities = softmax(scores)
|
||||
limit = min(top_k, probabilities.size)
|
||||
top_indices = np.argsort(probabilities)[-limit:][::-1]
|
||||
results = []
|
||||
for index in top_indices:
|
||||
label = None
|
||||
if labels is not None and 0 <= int(index) < len(labels):
|
||||
label = labels[int(index)]
|
||||
results.append(
|
||||
{
|
||||
"index": int(index),
|
||||
"label": label,
|
||||
"probability": float(probabilities[int(index)]),
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def render_result_line(result) -> str:
|
||||
name = result["label"] if result["label"] else f'class {result["index"]}'
|
||||
return f'{name}: {result["probability"] * 100.0:.2f}%'
|
||||
|
||||
|
||||
def draw_classification_panel(image: Image.Image, results, output_path: Path):
|
||||
annotated = image.copy()
|
||||
draw = ImageDraw.Draw(annotated)
|
||||
lines = [render_result_line(result) for result in results]
|
||||
if not lines:
|
||||
lines = ["No predictions"]
|
||||
|
||||
padding = 10
|
||||
line_gap = 4
|
||||
max_width = 0
|
||||
line_heights = []
|
||||
for line in lines:
|
||||
left, top, right, bottom = draw.textbbox((0, 0), line)
|
||||
max_width = max(max_width, right - left)
|
||||
line_heights.append(bottom - top)
|
||||
|
||||
panel_height = padding * 2 + sum(line_heights) + line_gap * (len(lines) - 1)
|
||||
panel_width = padding * 2 + max_width
|
||||
origin_x = 12
|
||||
origin_y = 12
|
||||
draw.rounded_rectangle(
|
||||
(origin_x, origin_y, origin_x + panel_width, origin_y + panel_height),
|
||||
radius=10,
|
||||
fill=(0, 0, 0),
|
||||
)
|
||||
|
||||
y = origin_y + padding
|
||||
for line, line_height in zip(lines, line_heights):
|
||||
draw.text((origin_x + padding, y), line, fill=(255, 255, 255))
|
||||
y += line_height + line_gap
|
||||
|
||||
annotated.save(output_path)
|
||||
|
||||
|
||||
def run_reference_and_simulator(args, model_path: Path, tensor: np.ndarray):
|
||||
model_dir = model_path.parent
|
||||
runner_build_dir = model_dir / "runner" / "build"
|
||||
runner_path = runner_build_dir / "runner"
|
||||
pim_dir = model_dir / "raptor" / "pim"
|
||||
simulation_dir = model_dir / "classification_demo" / "simulation"
|
||||
reference_dir = model_dir / "classification_demo" / "reference"
|
||||
inputs_dir = model_dir / "classification_demo" / "inputs"
|
||||
|
||||
simulation_dir.mkdir(parents=True, exist_ok=True)
|
||||
reference_dir.mkdir(parents=True, exist_ok=True)
|
||||
inputs_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
input_descriptors, output_descriptors = onnx_io(model_path)
|
||||
if len(input_descriptors) != 1:
|
||||
raise RuntimeError(f"Expected one classification input tensor, found {len(input_descriptors)}")
|
||||
if len(output_descriptors) != 1:
|
||||
raise RuntimeError(f"Expected one classification output tensor, found {len(output_descriptors)}")
|
||||
|
||||
input_index, _input_name, _input_dtype, input_shape = input_descriptors[0]
|
||||
if list(tensor.shape) != list(input_shape):
|
||||
raise RuntimeError(f"Preprocessed tensor shape {list(tensor.shape)} does not match model input {input_shape}")
|
||||
|
||||
input_csv = inputs_dir / "in0.csv"
|
||||
save_tensor_csv(tensor, input_csv)
|
||||
|
||||
runner_cmd = [
|
||||
str(runner_path),
|
||||
f"--in{input_index}-csv-file",
|
||||
str(input_csv),
|
||||
f"--in{input_index}-shape",
|
||||
"x".join(str(dim) for dim in tensor.shape),
|
||||
"--save-csv-dir",
|
||||
str(reference_dir),
|
||||
]
|
||||
subprocess.run(runner_cmd, cwd=runner_build_dir, check=True)
|
||||
|
||||
write_inputs_to_memory_bin(pim_dir / "memory.bin", pim_dir / "config.json", [tensor])
|
||||
dump_ranges = build_dump_ranges(pim_dir / "config.json", output_descriptors)
|
||||
output_bin_path = simulation_dir / "out.bin"
|
||||
run_pim_simulator(
|
||||
args.simulator_dir,
|
||||
pim_dir,
|
||||
output_bin_path,
|
||||
dump_ranges,
|
||||
timeout_sec=args.command_timeout_seconds,
|
||||
)
|
||||
|
||||
output_index, output_name, output_dtype_code, output_shape = output_descriptors[0]
|
||||
output_dtype = np.dtype(_ONNX_TO_NP[output_dtype_code])
|
||||
reference_csv = reference_dir / f"output{output_index}_{sanitize_output_name(output_name)}.csv"
|
||||
reference_output = np.loadtxt(reference_csv, delimiter=",", dtype=output_dtype).reshape(output_shape)
|
||||
simulator_output = parse_pim_simulator_outputs(output_bin_path, output_descriptors)[0]
|
||||
return reference_output, simulator_output
|
||||
|
||||
|
||||
def print_topk(title: str, results):
|
||||
print(title)
|
||||
for rank, result in enumerate(results, start=1):
|
||||
label_text = result["label"] if result["label"] else f'class {result["index"]}'
|
||||
print(f' {rank}. {label_text} ({result["probability"] * 100.0:.2f}%) [index={result["index"]}]')
|
||||
|
||||
|
||||
def main():
|
||||
defaults = resolve_default_paths()
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run a VGG or ResNet ONNX model through the Raptor simulator and annotate the image with top classification results.")
|
||||
parser.add_argument("--model", type=Path, default=None)
|
||||
parser.add_argument("--network", choices=("resnet", "vgg"), default=None)
|
||||
parser.add_argument("--image", type=Path, required=True)
|
||||
parser.add_argument("--labels", type=Path, default=None)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--raptor-path", type=Path, default=defaults["raptor_path"])
|
||||
parser.add_argument("--onnx-include-dir", type=Path, default=defaults["onnx_include_dir"])
|
||||
parser.add_argument("--simulator-dir", type=Path, default=defaults["simulator_dir"])
|
||||
parser.add_argument("--crossbar-size", type=int, default=2048)
|
||||
parser.add_argument("--crossbar-count", type=int, default=256)
|
||||
parser.add_argument("--core-count", type=int, default=1000)
|
||||
parser.add_argument("--top-k", type=int, default=5)
|
||||
parser.add_argument("--command-timeout-seconds", type=float, default=7200.0)
|
||||
parser.add_argument("--skip-compile", action="store_true")
|
||||
parser.add_argument("--verbose", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
args.model = resolve_model_path(args.network, args.model)
|
||||
args.image = args.image.resolve()
|
||||
args.output = args.output.resolve()
|
||||
args.labels = args.labels.resolve() if args.labels else None
|
||||
args.raptor_path = args.raptor_path.resolve()
|
||||
args.onnx_include_dir = args.onnx_include_dir.resolve()
|
||||
args.simulator_dir = args.simulator_dir.resolve()
|
||||
|
||||
if not args.skip_compile:
|
||||
ensure_local_artifacts(args, args.model)
|
||||
else:
|
||||
ensure_existing_artifacts(args.model.parent)
|
||||
|
||||
original_image, tensor = preprocess_classification_image(args.image)
|
||||
labels = load_labels(args.labels)
|
||||
reference_output, simulator_output = run_reference_and_simulator(args, args.model, tensor)
|
||||
reference_results = decode_classification_output(reference_output, labels, args.top_k)
|
||||
simulator_results = decode_classification_output(simulator_output, labels, args.top_k)
|
||||
|
||||
print_topk("Reference top-k:", reference_results)
|
||||
print_topk("Simulator top-k:", simulator_results)
|
||||
|
||||
reference_scores = np.asarray(reference_output, dtype=np.float64).reshape(-1)
|
||||
simulator_scores = np.asarray(simulator_output, dtype=np.float64).reshape(-1)
|
||||
max_abs_diff = float(np.max(np.abs(reference_scores - simulator_scores)))
|
||||
print(f"Max absolute score diff: {max_abs_diff:.6e}")
|
||||
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
draw_classification_panel(original_image, simulator_results, args.output)
|
||||
print(f"Annotated image saved to {args.output}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user