nvidia

/

nemotron-graphic-elements-v1

like 13

Follow

NVIDIA 44.9k

(hf_nvingest_models) vjawa@dgx-a100-01:/raid/vjawa/hf_models/nemotron-table-structure-v1$ cat preprocess_benchmark.py
#!/usr/bin/env python3
"""Benchmark the preprocessing step for different batch sizes."""

i#!/usr/bin/env python3
"""Benchmark the preprocessing step for different batch sizes."""

import argparse
import time
from typing import Iterable, List

import numpy as np
import torch
from PIL import Image

from nemotron_graphic_elements_v1.model import define_model


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Benchmark model.preprocess for various batch sizes."
    )
    parser.add_argument(
        "--image",
        default="./example.png",
        help="Path to an RGB image to use for benchmarking.",
    )
    parser.add_argument(
        "--batch-sizes",
        type=int,
        nargs="+",
        default=[1, 2, 4, 8, 16, 32],
        help="Batch sizes to benchmark.",
    )
    parser.add_argument(
        "--warmup-iters",
        type=int,
        default=3,
        help="Number of warmup iterations to run before measuring.",
    )
    parser.add_argument(
        "--measure-iters",
        type=int,
        default=10,
        help="Number of measurement iterations per batch size.",
    )
    return parser.parse_args()


def load_image(path: str) -> np.ndarray:
    """Load an RGB image as a numpy array."""
    return np.array(Image.open(path).convert("RGB"))


def benchmark_preprocess(
    model: torch.nn.Module,
    image: np.ndarray,
    batch_sizes: Iterable[int],
    warmup_iters: int,
    measure_iters: int,
) -> List[dict]:
    results: List[dict] = []
    for bs in batch_sizes:
        # Reuse the same image; preprocess does not mutate the input.
        images = [image.copy() for _ in range(bs)]

        # Warmup
        with torch.inference_mode():
            for _ in range(warmup_iters):
                for img in images:
                    output = model.preprocess(img)

        print("Preprocess output device: ", output.device)
        # Measure
        times_s: List[float] = []
        for _ in range(measure_iters):
            start = time.perf_counter()
            with torch.inference_mode():
                batch_tensors = [model.preprocess(img) for img in images]
                _ = torch.stack(batch_tensors)
            end = time.perf_counter()
            times_s.append(end - start)

        times = np.array(times_s)
        mean_ms = float(times.mean() * 1000)
        p50_ms = float(np.percentile(times, 50) * 1000)
        p95_ms = float(np.percentile(times, 95) * 1000)
        imgs_per_sec = float(bs / times.mean())

        results.append(
            {
                "batch_size": bs,
                "mean_ms": mean_ms,
                "p50_ms": p50_ms,
                "p95_ms": p95_ms,
                "imgs_per_sec": imgs_per_sec,
            }
        )

    return results


def main() -> None:
    args = parse_args()

    image = load_image(args.image)
    model = define_model("graphic_element_v1", verbose=False)

    results = benchmark_preprocess(
        model,
        image,
        args.batch_sizes,
        args.warmup_iters,
        args.measure_iters,
    )

    print(
        f"Benchmarking preprocess on device: {model.device}, "
        f"warmup={args.warmup_iters}, iters={args.measure_iters}"
    )
    header = f"{'BS':>4} | {'mean (ms)':>10} | {'p50 (ms)':>9} | {'p95 (ms)':>9} | {'imgs/s':>8}"
    print(header)
    print("-" * len(header))
    for res in results:
        print(
            f"{res['batch_size']:>4} | "
            f"{res['mean_ms']:10.2f} | "
            f"{res['p50_ms']:9.2f} | "
            f"{res['p95_ms']:9.2f} | "
            f"{res['imgs_per_sec']:8.2f}"
        )


if __name__ == "__main__":
    main()