| 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() |
