File size: 1,955 Bytes
9b92c75 | 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 | from __future__ import annotations
import argparse
import json
import time
import torch
from .config import apply_overrides, load_config
from .model import build_model
from .utils import trainable_parameter_count
def main() -> None:
parser = argparse.ArgumentParser(description="Profile ObjectModel-v1 parameters and latency")
parser.add_argument("--config", default="configs/objectmodel_v1.yaml")
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
parser.add_argument("--warmup", type=int, default=10)
parser.add_argument("--runs", type=int, default=50)
parser.add_argument("--set", action="append", default=[])
args = parser.parse_args()
config = apply_overrides(load_config(args.config), args.set)
device = torch.device(args.device)
model = build_model(config).eval().to(device)
size = model.spec.input_size
sample = torch.randn(1, 3, size, size, device=device)
with torch.inference_mode():
for _ in range(args.warmup):
model(sample)
if device.type == "cuda":
torch.cuda.synchronize()
durations = []
for _ in range(args.runs):
start = time.perf_counter()
model(sample)
if device.type == "cuda":
torch.cuda.synchronize()
durations.append((time.perf_counter() - start) * 1000)
durations.sort()
parameters = trainable_parameter_count(model)
report = {
"parameters": parameters,
"parameters_millions": round(parameters / 1e6, 3),
"fp32_weight_megabytes": round(parameters * 4 / 1024**2, 2),
"input_size": size,
"device": str(device),
"latency_ms_median": round(durations[len(durations) // 2], 3),
"latency_ms_p95": round(durations[min(int(len(durations) * 0.95), len(durations) - 1)], 3),
}
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main() |