File size: 2,114 Bytes
299866a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from __future__ import annotations
import argparse
from pathlib import Path
from PIL import Image
from app.config import get_settings
from app.model_backend import (
benchmark_predictions,
get_onnx_classifier,
get_torch_classifier,
serialize_benchmark_report,
)
def load_sample_image(image_path: Path) -> bytes:
with Image.open(image_path) as image:
rgb_image = image.convert("RGB")
from io import BytesIO
buffer = BytesIO()
rgb_image.save(buffer, format="JPEG")
return buffer.getvalue()
def model_size_mb(model_path: Path) -> float:
return model_path.stat().st_size / (1024 * 1024)
def main() -> None:
parser = argparse.ArgumentParser(description="Benchmark original, ONNX, and quantized model variants.")
parser.add_argument("--image", type=Path, required=True, help="Sample image used for benchmarking")
args = parser.parse_args()
settings = get_settings()
image_bytes = load_sample_image(args.image)
torch_backend = get_torch_classifier(settings.hf_model_name)
onnx_backend = get_onnx_classifier(str(settings.onnx_path), settings.hf_model_name)
quantized_backend = get_onnx_classifier(str(settings.quantized_onnx_path), settings.hf_model_name)
rows = [
{
"variant": "Original",
"model_size_mb": model_size_mb(settings.torch_weights_path),
**benchmark_predictions(torch_backend.predict, image_bytes),
},
{
"variant": "ONNX",
"model_size_mb": model_size_mb(settings.onnx_path),
**benchmark_predictions(onnx_backend.predict, image_bytes),
},
{
"variant": "Quantized",
"model_size_mb": model_size_mb(settings.quantized_onnx_path),
**benchmark_predictions(quantized_backend.predict, image_bytes),
},
]
settings.docs_dir.mkdir(parents=True, exist_ok=True)
settings.benchmark_output_path.write_text(serialize_benchmark_report(rows), encoding="utf-8")
print(serialize_benchmark_report(rows))
if __name__ == "__main__":
main()
|