torchmeter-cli 0.1.0__py3-none-any.whl

torchmeter_cli/__init__.py

@@ -0,0 +1,11 @@
+"""TorchMeter CLI: PyTorch model performance benchmarking."""
+
+from .benchmark import benchmark
+from .result import BenchmarkResult
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "BenchmarkResult",
+    "benchmark",
+]

torchmeter_cli/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

torchmeter_cli/benchmark.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import statistics
+import time
+from collections.abc import Callable
+
+import torch
+
+from .result import BenchmarkResult
+
+
+def _synchronize(device: torch.device) -> None:
+    """Synchronize CUDA so GPU operations are fully completed."""
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+
+
+def _percentile(values: list[float], percentile: float) -> float:
+    """Calculate a percentile using linear interpolation."""
+    if not values:
+        raise ValueError("Cannot calculate percentile from empty data.")
+
+    values = sorted(values)
+
+    if len(values) == 1:
+        return values[0]
+
+    index = (len(values) - 1) * percentile
+    lower = int(index)
+    upper = min(lower + 1, len(values) - 1)
+
+    weight = index - lower
+
+    return values[lower] + (values[upper] - values[lower]) * weight
+
+def _set_batch_size(inputs: object, batch_size: int) -> object:
+    """Create inputs with the requested batch size."""
+    if torch.is_tensor(inputs):
+        if inputs.ndim == 0:
+            return inputs
+
+        shape = list(inputs.shape)
+        shape[0] = batch_size
+
+        return torch.randn(
+            shape,
+            dtype=inputs.dtype,
+            device=inputs.device,
+        )
+
+    if isinstance(inputs, dict):
+        return {
+            key: _set_batch_size(value, batch_size)
+            for key, value in inputs.items()
+        }
+
+    if isinstance(inputs, tuple):
+        return tuple(
+            _set_batch_size(value, batch_size)
+            for value in inputs
+        )
+
+    if isinstance(inputs, list):
+        return [
+            _set_batch_size(value, batch_size)
+            for value in inputs
+        ]
+
+    return inputs
+def _run_model(
+    model: torch.nn.Module,
+    inputs: object,
+) -> object:
+    """Run the model with the provided inputs."""
+    if isinstance(inputs, dict):
+        return model(**inputs)
+
+    if isinstance(inputs, (tuple, list)):
+        return model(*inputs)
+
+    return model(inputs)
+
+
+def _move_to_device(
+    value: object,
+    device: torch.device,
+) -> object:
+    """Move tensors or nested structures to the target device."""
+    if torch.is_tensor(value):
+        return value.to(device)
+
+    if isinstance(value, dict):
+        return {
+            key: _move_to_device(item, device)
+            for key, item in value.items()
+        }
+
+    if isinstance(value, tuple):
+        return tuple(_move_to_device(item, device) for item in value)
+
+    if isinstance(value, list):
+        return [_move_to_device(item, device) for item in value]
+
+    return value
+
+
+def benchmark(
+    model: torch.nn.Module,
+    inputs: object,
+    *,
+    model_name: str = "PyTorch Model",
+    device: str = "auto",
+    warmup: int = 20,
+    iterations: int = 100,
+) -> BenchmarkResult:
+    """Benchmark a PyTorch model.
+
+    Args:
+        model: PyTorch model to benchmark.
+        inputs: Example model input.
+        model_name: Human-readable model name.
+        device: "auto", "cpu", or "cuda".
+        warmup: Number of warmup iterations.
+        iterations: Number of measured iterations.
+    """
+    if warmup < 0:
+        raise ValueError("warmup must be >= 0")
+
+    if iterations <= 0:
+        raise ValueError("iterations must be > 0")
+
+    if device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    target_device = torch.device(device)
+
+    if target_device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA was requested but is not available.")
+
+    model = model.to(target_device)
+    model.eval()
+
+    inputs = _move_to_device(inputs, target_device)
+
+    if target_device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(target_device)
+
+    with torch.inference_mode():
+        for _ in range(warmup):
+            _run_model(model, inputs)
+
+        _synchronize(target_device)
+
+        timings: list[float] = []
+
+        for _ in range(iterations):
+            _synchronize(target_device)
+
+            start = time.perf_counter()
+
+            _run_model(model, inputs)
+
+            _synchronize(target_device)
+
+            end = time.perf_counter()
+
+            timings.append((end - start) * 1000)
+
+    mean_latency = statistics.mean(timings)
+    median_latency = statistics.median(timings)
+
+    p95 = _percentile(timings, 0.95)
+    p99 = _percentile(timings, 0.99)
+
+    throughput = 1000.0 / mean_latency
+
+    peak_memory = None
+
+    if target_device.type == "cuda":
+        peak_memory = (
+            torch.cuda.max_memory_allocated(target_device)
+            / (1024**2)
+        )
+
+    return BenchmarkResult(
+        model_name=model_name,
+        device=str(target_device),
+        iterations=iterations,
+        warmup_iterations=warmup,
+        mean_latency_ms=mean_latency,
+        median_latency_ms=median_latency,
+        p95_latency_ms=p95,
+        p99_latency_ms=p99,
+        throughput_samples_per_second=throughput,
+        peak_memory_mb=peak_memory,
+    )

torchmeter_cli/cli.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import sys
+from pathlib import Path
+
+from .benchmark import benchmark
+
+
+def _load_module(path: Path):
+    """Load a Python file as a module."""
+    spec = importlib.util.spec_from_file_location(
+        "torchmeter_cli_user_model",
+        path,
+    )
+
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Could not load {path}")
+
+    module = importlib.util.module_from_spec(spec)
+
+    sys.path.insert(0, str(path.parent))
+
+    try:
+        spec.loader.exec_module(module)
+    finally:
+        sys.path.pop(0)
+
+    return module
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="torchmeter-cli",
+        description="Performance benchmarking for PyTorch models.",
+    )
+
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    run_parser = subparsers.add_parser(
+        "run",
+        help="Benchmark a Python model file.",
+    )
+
+    run_parser.add_argument(
+        "model",
+        type=Path,
+        help="Path to the model.py file.",
+    )
+
+    run_parser.add_argument(
+        "--device",
+        choices=["auto", "cpu", "cuda"],
+        default="auto",
+    )
+
+    run_parser.add_argument(
+        "--warmup",
+        type=int,
+        default=20,
+    )
+
+    run_parser.add_argument(
+        "--iterations",
+        type=int,
+        default=100,
+    )
+
+    run_parser.add_argument(
+        "--format",
+        choices=["text", "json", "markdown"],
+        default="text",
+    )
+
+    run_parser.add_argument(
+        "--output",
+        type=Path,
+        help="Write the result to a file.",
+    )
+    run_parser.add_argument(
+        "--batch-sizes",
+        type=str,
+        help="Comma-separated batch sizes to benchmark, e.g. 1,8,16,32.",
+    )
+    return parser
+
+
+def _print_text(result) -> None:
+    print()
+    print("TorchMeter Benchmark")
+    print("─" * 32)
+    print(f"Model:       {result.model_name}")
+    print(f"Device:      {result.device}")
+    print(f"Iterations:  {result.iterations}")
+    print()
+    print("Latency")
+    print(f"  Mean:       {result.mean_latency_ms:.3f} ms")
+    print(f"  Median:     {result.median_latency_ms:.3f} ms")
+    print(f"  P95:        {result.p95_latency_ms:.3f} ms")
+    print(f"  P99:        {result.p99_latency_ms:.3f} ms")
+    print()
+    print(f"Throughput:  {result.throughput_samples_per_second:.2f} samples/sec")
+
+    if result.peak_memory_mb is not None:
+        print(f"Peak memory: {result.peak_memory_mb:.2f} MB")
+
+    print()
+
+
+def main() -> None:
+    parser = _build_parser()
+    args = parser.parse_args()
+
+    if args.command != "run":
+        parser.print_help()
+        return
+
+    if not args.model.exists():
+        parser.error(f"Model file does not exist: {args.model}")
+
+    try:
+        module = _load_module(args.model)
+
+        if not hasattr(module, "create_model"):
+            parser.error("model.py must define create_model().")
+
+        if not hasattr(module, "create_input"):
+            parser.error("model.py must define create_input().")
+
+        model = module.create_model()
+        inputs = module.create_input()
+
+        model_name = getattr(
+            module,
+            "MODEL_NAME",
+            model.__class__.__name__,
+        )
+
+        result = benchmark(
+            model,
+            inputs,
+            model_name=model_name,
+            device=args.device,
+            warmup=args.warmup,
+            iterations=args.iterations,
+        )
+
+    except Exception as exc:
+        parser.error(str(exc))
+
+    if args.format == "json":
+        output = result.to_json()
+    elif args.format == "markdown":
+        output = result.to_markdown()
+    else:
+        _print_text(result)
+        output = None
+
+    if output is not None:
+        if args.output:
+            args.output.write_text(output + "\n")
+            print(f"Report written to {args.output}")
+        else:
+            print(output)
+
+
+if __name__ == "__main__":
+    main()

torchmeter_cli/result.py

@@ -0,0 +1,47 @@
+from dataclasses import asdict, dataclass
+import json
+
+
+@dataclass
+class BenchmarkResult:
+    model_name: str
+    device: str
+    iterations: int
+    warmup_iterations: int
+
+    mean_latency_ms: float
+    median_latency_ms: float
+    p95_latency_ms: float
+    p99_latency_ms: float
+
+    throughput_samples_per_second: float
+    peak_memory_mb: float | None = None
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict(), indent=2)
+
+    def to_markdown(self) -> str:
+        memory = (
+            f"{self.peak_memory_mb:.2f} MB"
+            if self.peak_memory_mb is not None
+            else "N/A"
+        )
+
+        return f"""# TorchMeter Benchmark
+
+| Metric | Value |
+|---|---:|
+| Model | `{self.model_name}` |
+| Device | `{self.device}` |
+| Iterations | {self.iterations} |
+| Warmup | {self.warmup_iterations} |
+| Mean latency | {self.mean_latency_ms:.3f} ms |
+| Median latency | {self.median_latency_ms:.3f} ms |
+| P95 latency | {self.p95_latency_ms:.3f} ms |
+| P99 latency | {self.p99_latency_ms:.3f} ms |
+| Throughput | {self.throughput_samples_per_second:.2f} samples/sec |
+| Peak GPU memory | {memory} |
+"""

torchmeter_cli-0.1.0.dist-info/METADATA

@@ -0,0 +1,425 @@
+Metadata-Version: 2.3
+Name: torchmeter-cli
+Version: 0.1.0
+Summary: Simple, reproducible performance benchmarking for PyTorch models.
+Author: Pujan Neupane
+Author-email: Pujan Neupane <pujanneupaneop0907@gmail.com>
+License: Apache-2.0
+Requires-Dist: torch>=2.0
+Requires-Dist: torchvision>=0.15
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# TorchMeter
+
+Simple, reproducible performance benchmarking for PyTorch models.
+
+TorchMeter makes it easy to measure the performance of PyTorch models and generate useful benchmark results from the command line.
+
+## Status
+
+🚧 **Early Development — v0.1.0**
+
+TorchMeter is currently an experimental open-source project.
+
+The API and CLI may change as the project develops.
+
+## Features
+
+- PyTorch model benchmarking
+- CPU benchmarking
+- CUDA benchmarking
+- Warmup iterations
+- Configurable benchmark iterations
+- Mean latency
+- Median latency
+- P95 latency
+- P99 latency
+- Throughput measurement
+- Peak CUDA memory measurement
+- JSON output
+- Markdown reports
+- Simple Python API
+- Reproducible environment information
+
+## Installation
+
+### Using uv
+
+```bash
+uv add torchmeter-cli
+```
+
+### From source
+
+Clone the repository:
+
+```bash
+git clone https://github.com/YOUR_USERNAME/torchmeter-cli.git
+cd torchmeter-cli
+```
+
+Install the development environment:
+
+```bash
+uv sync
+```
+
+## Quick Start
+
+Create a Python file containing a model and example input.
+
+For example:
+
+```python
+import torch
+import torchvision.models as models
+
+
+MODEL_NAME = "ResNet50"
+
+
+def create_model():
+    return models.resnet50(weights=None)
+
+
+def create_input():
+    return torch.randn(1, 3, 224, 224)
+```
+
+Save it as:
+
+```text
+examples/model.py
+```
+
+Run the benchmark:
+
+```bash
+uv run torchmeter-cli run examples/model.py
+```
+
+For CUDA:
+
+```bash
+uv run torchmeter-cli run examples/model.py --device cuda
+```
+
+For CPU:
+
+```bash
+uv run torchmeter-cli run examples/model.py --device cpu
+```
+
+## Model Interface
+
+TorchMeter expects the model file to provide two functions:
+
+```python
+def create_model():
+    ...
+
+
+def create_input():
+    ...
+```
+
+You can optionally provide a model name:
+
+```python
+MODEL_NAME = "My Model"
+```
+
+If `MODEL_NAME` is not provided, TorchMeter uses the model class name.
+
+## Example Models
+
+The repository includes several example models:
+
+```text
+examples/
+├── cnn.py
+├── lstm.py
+├── mlp.py
+├── model.py
+├── rnn.py
+├── resnet.py
+└── transformer.py
+```
+
+Run an example:
+
+```bash
+uv run torchmeter-cli run examples/cnn.py --device cuda
+```
+
+```bash
+uv run torchmeter-cli run examples/rnn.py --device cuda
+```
+
+```bash
+uv run torchmeter-cli run examples/lstm.py --device cuda
+```
+
+```bash
+uv run torchmeter-cli run examples/transformer.py --device cuda
+```
+
+```bash
+uv run torchmeter-cli run examples/resnet.py --device cuda
+```
+
+## Benchmark Configuration
+
+Control the number of warmup iterations:
+
+```bash
+uv run torchmeter-cli run examples/model.py --warmup 20
+```
+
+Control the number of benchmark iterations:
+
+```bash
+uv run torchmeter-cli run examples/model.py --iterations 100
+```
+
+You can combine options:
+
+```bash
+uv run torchmeter-cli run examples/model.py \
+    --device cuda \
+    --warmup 20 \
+    --iterations 100
+```
+
+## Output Formats
+
+### Terminal
+
+The default output is a human-readable benchmark:
+
+```text
+TorchMeter Benchmark
+────────────────────────────────
+Model:       ResNet50
+Device:      cuda
+Iterations:  100
+
+Latency
+  Mean:       4.823 ms
+  Median:     4.761 ms
+  P95:        5.132 ms
+  P99:        5.421 ms
+
+Throughput:  207.34 samples/sec
+Peak memory: 1420.31 MB
+```
+
+### JSON
+
+Generate machine-readable output:
+
+```bash
+uv run torchmeter-cli run examples/model.py \
+    --format json
+```
+
+Write the JSON result to a file:
+
+```bash
+uv run torchmeter-cli run examples/model.py \
+    --format json \
+    --output results.json
+```
+
+### Markdown
+
+Generate a Markdown report:
+
+```bash
+uv run torchmeter-cli run examples/model.py \
+    --format markdown
+```
+
+Write it to a file:
+
+```bash
+uv run torchmeter-cli run examples/model.py \
+    --format markdown \
+    --output benchmark.md
+```
+
+## Python API
+
+TorchMeter can also be used directly from Python.
+
+```python
+import torch
+
+from torchmeter_cli import benchmark
+
+
+model = torch.nn.Sequential(
+    torch.nn.Linear(768, 1024),
+    torch.nn.ReLU(),
+    torch.nn.Linear(1024, 10),
+)
+
+inputs = torch.randn(1, 768)
+
+result = benchmark(
+    model,
+    inputs,
+    device="cuda",
+    warmup=20,
+    iterations=100,
+)
+
+print(result.mean_latency_ms)
+print(result.throughput_samples_per_second)
+```
+
+## Benchmark Results
+
+A benchmark contains:
+
+- Model name
+- Device
+- Number of warmup iterations
+- Number of benchmark iterations
+- Mean latency
+- Median latency
+- P95 latency
+- P99 latency
+- Throughput
+- Peak CUDA memory when available
+
+## Development
+
+Install dependencies:
+
+```bash
+uv sync
+```
+
+Run tests:
+
+```bash
+uv run pytest
+```
+
+Run the test suite with verbose output:
+
+```bash
+uv run pytest -v
+```
+
+Run Ruff:
+
+```bash
+uv run ruff check .
+```
+
+Format the project:
+
+```bash
+uv run ruff format .
+```
+
+## Project Structure
+
+```text
+torchmeter-cli/
+├── examples/
+│   ├── cnn.py
+│   ├── lstm.py
+│   ├── mlp.py
+│   ├── model.py
+│   ├── rnn.py
+│   ├── resnet.py
+│   └── transformer.py
+│
+├── src/
+│   └── torchmeter-cli/
+│       ├── __init__.py
+│       ├── __main__.py
+│       ├── benchmark.py
+│       ├── cli.py
+│       └── result.py
+│
+├── tests/
+│   └── test_benchmark.py
+│
+├── LICENSE
+├── README.md
+├── pyproject.toml
+└── uv.lock
+```
+
+## Design Goals
+
+TorchMeter aims to be:
+
+1. **Simple** — benchmark a model with a single command.
+2. **Reproducible** — capture the environment used for benchmarking.
+3. **Accurate** — correctly measure CPU and CUDA workloads.
+4. **Scriptable** — provide JSON output for automation.
+5. **Extensible** — provide both a CLI and Python API.
+6. **Open source** — developed transparently with community contributions.
+
+## Roadmap
+
+### v0.1.x
+
+- [x] Basic PyTorch benchmarking
+- [x] CPU support
+- [x] CUDA support
+- [x] Latency statistics
+- [x] Throughput measurement
+- [x] CUDA memory measurement
+- [x] JSON output
+- [x] Markdown output
+- [x] Python API
+- [x] Example models
+- [x] Basic tests
+
+### Future
+
+Potential future features include:
+
+- `torch.compile` comparisons
+- Performance regression detection
+- Benchmark baselines
+- CI integration
+- GitHub Actions
+- Additional hardware metrics
+- Distributed benchmarking
+- Benchmark history
+- Performance dashboards
+
+Future features are subject to change based on project feedback and development priorities.
+
+## Contributing
+
+Contributions are welcome.
+
+Before opening a pull request, please run:
+
+```bash
+uv run pytest
+uv run ruff check .
+uv run ruff format --check .
+```
+
+For larger changes, please open an issue first to discuss the proposed approach.
+
+## License
+
+TorchMeter is licensed under the Apache License 2.0.
+
+See [LICENSE](LICENSE) for the full license text.
+
+## Disclaimer
+
+TorchMeter is an early-stage project. Benchmark results can vary depending on hardware, software versions, system load, and benchmark configuration.
+Always compare results under controlled conditions when making performance decisions.

torchmeter_cli-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+torchmeter_cli/__init__.py,sha256=CX8yNMm1XoEpvB47uYqXVxEljBu_jwrNEkIw6xeyb2k,209
+torchmeter_cli/__main__.py,sha256=MSmt_5Xg84uHqzTN38JwgseJK8rsJn_11A8WD99VtEo,61
+torchmeter_cli/benchmark.py,sha256=4-w2eRtDkR8Rr981AeCZ0ftRuypqVj2rN4hWJHKDIOo,4963
+torchmeter_cli/cli.py,sha256=OinZK5qzJ-7_7rC6xu87-TTwTajUh2x4ybl9uhK84Dc,4016
+torchmeter_cli/result.py,sha256=BQvh1ZVi5WzSXUVMPECDHFScL6I3EPF9fMVTyGX22gk,1194
+torchmeter_cli-0.1.0.dist-info/WHEEL,sha256=lrO5MD1WVAWzcbNy_L2BwtfPrcM3KfUFpbKYXLkJX4A,80
+torchmeter_cli-0.1.0.dist-info/entry_points.txt,sha256=AUVirtlEcDGR49Li0xaB4W1BEUW0KCC9fGNV5JlMM88,60
+torchmeter_cli-0.1.0.dist-info/METADATA,sha256=J6WWsqa1oOQBRtVDb2Jvk1YWJdjIpCdV4_BA8VMhKR0,7411
+torchmeter_cli-0.1.0.dist-info/RECORD,,

torchmeter_cli-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.1
+Root-Is-Purelib: true
+Tag: py3-none-any

torchmeter_cli-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+torchmeter-cli = torchmeter_cli.cli:main
+