torchmeter-cli 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,425 @@
1
+ Metadata-Version: 2.3
2
+ Name: torchmeter-cli
3
+ Version: 0.1.0
4
+ Summary: Simple, reproducible performance benchmarking for PyTorch models.
5
+ Author: Pujan Neupane
6
+ Author-email: Pujan Neupane <pujanneupaneop0907@gmail.com>
7
+ License: Apache-2.0
8
+ Requires-Dist: torch>=2.0
9
+ Requires-Dist: torchvision>=0.15
10
+ Requires-Python: >=3.11
11
+ Description-Content-Type: text/markdown
12
+
13
+ # TorchMeter
14
+
15
+ Simple, reproducible performance benchmarking for PyTorch models.
16
+
17
+ TorchMeter makes it easy to measure the performance of PyTorch models and generate useful benchmark results from the command line.
18
+
19
+ ## Status
20
+
21
+ 🚧 **Early Development — v0.1.0**
22
+
23
+ TorchMeter is currently an experimental open-source project.
24
+
25
+ The API and CLI may change as the project develops.
26
+
27
+ ## Features
28
+
29
+ - PyTorch model benchmarking
30
+ - CPU benchmarking
31
+ - CUDA benchmarking
32
+ - Warmup iterations
33
+ - Configurable benchmark iterations
34
+ - Mean latency
35
+ - Median latency
36
+ - P95 latency
37
+ - P99 latency
38
+ - Throughput measurement
39
+ - Peak CUDA memory measurement
40
+ - JSON output
41
+ - Markdown reports
42
+ - Simple Python API
43
+ - Reproducible environment information
44
+
45
+ ## Installation
46
+
47
+ ### Using uv
48
+
49
+ ```bash
50
+ uv add torchmeter-cli
51
+ ```
52
+
53
+ ### From source
54
+
55
+ Clone the repository:
56
+
57
+ ```bash
58
+ git clone https://github.com/YOUR_USERNAME/torchmeter-cli.git
59
+ cd torchmeter-cli
60
+ ```
61
+
62
+ Install the development environment:
63
+
64
+ ```bash
65
+ uv sync
66
+ ```
67
+
68
+ ## Quick Start
69
+
70
+ Create a Python file containing a model and example input.
71
+
72
+ For example:
73
+
74
+ ```python
75
+ import torch
76
+ import torchvision.models as models
77
+
78
+
79
+ MODEL_NAME = "ResNet50"
80
+
81
+
82
+ def create_model():
83
+ return models.resnet50(weights=None)
84
+
85
+
86
+ def create_input():
87
+ return torch.randn(1, 3, 224, 224)
88
+ ```
89
+
90
+ Save it as:
91
+
92
+ ```text
93
+ examples/model.py
94
+ ```
95
+
96
+ Run the benchmark:
97
+
98
+ ```bash
99
+ uv run torchmeter-cli run examples/model.py
100
+ ```
101
+
102
+ For CUDA:
103
+
104
+ ```bash
105
+ uv run torchmeter-cli run examples/model.py --device cuda
106
+ ```
107
+
108
+ For CPU:
109
+
110
+ ```bash
111
+ uv run torchmeter-cli run examples/model.py --device cpu
112
+ ```
113
+
114
+ ## Model Interface
115
+
116
+ TorchMeter expects the model file to provide two functions:
117
+
118
+ ```python
119
+ def create_model():
120
+ ...
121
+
122
+
123
+ def create_input():
124
+ ...
125
+ ```
126
+
127
+ You can optionally provide a model name:
128
+
129
+ ```python
130
+ MODEL_NAME = "My Model"
131
+ ```
132
+
133
+ If `MODEL_NAME` is not provided, TorchMeter uses the model class name.
134
+
135
+ ## Example Models
136
+
137
+ The repository includes several example models:
138
+
139
+ ```text
140
+ examples/
141
+ ├── cnn.py
142
+ ├── lstm.py
143
+ ├── mlp.py
144
+ ├── model.py
145
+ ├── rnn.py
146
+ ├── resnet.py
147
+ └── transformer.py
148
+ ```
149
+
150
+ Run an example:
151
+
152
+ ```bash
153
+ uv run torchmeter-cli run examples/cnn.py --device cuda
154
+ ```
155
+
156
+ ```bash
157
+ uv run torchmeter-cli run examples/rnn.py --device cuda
158
+ ```
159
+
160
+ ```bash
161
+ uv run torchmeter-cli run examples/lstm.py --device cuda
162
+ ```
163
+
164
+ ```bash
165
+ uv run torchmeter-cli run examples/transformer.py --device cuda
166
+ ```
167
+
168
+ ```bash
169
+ uv run torchmeter-cli run examples/resnet.py --device cuda
170
+ ```
171
+
172
+ ## Benchmark Configuration
173
+
174
+ Control the number of warmup iterations:
175
+
176
+ ```bash
177
+ uv run torchmeter-cli run examples/model.py --warmup 20
178
+ ```
179
+
180
+ Control the number of benchmark iterations:
181
+
182
+ ```bash
183
+ uv run torchmeter-cli run examples/model.py --iterations 100
184
+ ```
185
+
186
+ You can combine options:
187
+
188
+ ```bash
189
+ uv run torchmeter-cli run examples/model.py \
190
+ --device cuda \
191
+ --warmup 20 \
192
+ --iterations 100
193
+ ```
194
+
195
+ ## Output Formats
196
+
197
+ ### Terminal
198
+
199
+ The default output is a human-readable benchmark:
200
+
201
+ ```text
202
+ TorchMeter Benchmark
203
+ ────────────────────────────────
204
+ Model: ResNet50
205
+ Device: cuda
206
+ Iterations: 100
207
+
208
+ Latency
209
+ Mean: 4.823 ms
210
+ Median: 4.761 ms
211
+ P95: 5.132 ms
212
+ P99: 5.421 ms
213
+
214
+ Throughput: 207.34 samples/sec
215
+ Peak memory: 1420.31 MB
216
+ ```
217
+
218
+ ### JSON
219
+
220
+ Generate machine-readable output:
221
+
222
+ ```bash
223
+ uv run torchmeter-cli run examples/model.py \
224
+ --format json
225
+ ```
226
+
227
+ Write the JSON result to a file:
228
+
229
+ ```bash
230
+ uv run torchmeter-cli run examples/model.py \
231
+ --format json \
232
+ --output results.json
233
+ ```
234
+
235
+ ### Markdown
236
+
237
+ Generate a Markdown report:
238
+
239
+ ```bash
240
+ uv run torchmeter-cli run examples/model.py \
241
+ --format markdown
242
+ ```
243
+
244
+ Write it to a file:
245
+
246
+ ```bash
247
+ uv run torchmeter-cli run examples/model.py \
248
+ --format markdown \
249
+ --output benchmark.md
250
+ ```
251
+
252
+ ## Python API
253
+
254
+ TorchMeter can also be used directly from Python.
255
+
256
+ ```python
257
+ import torch
258
+
259
+ from torchmeter_cli import benchmark
260
+
261
+
262
+ model = torch.nn.Sequential(
263
+ torch.nn.Linear(768, 1024),
264
+ torch.nn.ReLU(),
265
+ torch.nn.Linear(1024, 10),
266
+ )
267
+
268
+ inputs = torch.randn(1, 768)
269
+
270
+ result = benchmark(
271
+ model,
272
+ inputs,
273
+ device="cuda",
274
+ warmup=20,
275
+ iterations=100,
276
+ )
277
+
278
+ print(result.mean_latency_ms)
279
+ print(result.throughput_samples_per_second)
280
+ ```
281
+
282
+ ## Benchmark Results
283
+
284
+ A benchmark contains:
285
+
286
+ - Model name
287
+ - Device
288
+ - Number of warmup iterations
289
+ - Number of benchmark iterations
290
+ - Mean latency
291
+ - Median latency
292
+ - P95 latency
293
+ - P99 latency
294
+ - Throughput
295
+ - Peak CUDA memory when available
296
+
297
+ ## Development
298
+
299
+ Install dependencies:
300
+
301
+ ```bash
302
+ uv sync
303
+ ```
304
+
305
+ Run tests:
306
+
307
+ ```bash
308
+ uv run pytest
309
+ ```
310
+
311
+ Run the test suite with verbose output:
312
+
313
+ ```bash
314
+ uv run pytest -v
315
+ ```
316
+
317
+ Run Ruff:
318
+
319
+ ```bash
320
+ uv run ruff check .
321
+ ```
322
+
323
+ Format the project:
324
+
325
+ ```bash
326
+ uv run ruff format .
327
+ ```
328
+
329
+ ## Project Structure
330
+
331
+ ```text
332
+ torchmeter-cli/
333
+ ├── examples/
334
+ │ ├── cnn.py
335
+ │ ├── lstm.py
336
+ │ ├── mlp.py
337
+ │ ├── model.py
338
+ │ ├── rnn.py
339
+ │ ├── resnet.py
340
+ │ └── transformer.py
341
+ │
342
+ ├── src/
343
+ │ └── torchmeter-cli/
344
+ │ ├── __init__.py
345
+ │ ├── __main__.py
346
+ │ ├── benchmark.py
347
+ │ ├── cli.py
348
+ │ └── result.py
349
+ │
350
+ ├── tests/
351
+ │ └── test_benchmark.py
352
+ │
353
+ ├── LICENSE
354
+ ├── README.md
355
+ ├── pyproject.toml
356
+ └── uv.lock
357
+ ```
358
+
359
+ ## Design Goals
360
+
361
+ TorchMeter aims to be:
362
+
363
+ 1. **Simple** — benchmark a model with a single command.
364
+ 2. **Reproducible** — capture the environment used for benchmarking.
365
+ 3. **Accurate** — correctly measure CPU and CUDA workloads.
366
+ 4. **Scriptable** — provide JSON output for automation.
367
+ 5. **Extensible** — provide both a CLI and Python API.
368
+ 6. **Open source** — developed transparently with community contributions.
369
+
370
+ ## Roadmap
371
+
372
+ ### v0.1.x
373
+
374
+ - [x] Basic PyTorch benchmarking
375
+ - [x] CPU support
376
+ - [x] CUDA support
377
+ - [x] Latency statistics
378
+ - [x] Throughput measurement
379
+ - [x] CUDA memory measurement
380
+ - [x] JSON output
381
+ - [x] Markdown output
382
+ - [x] Python API
383
+ - [x] Example models
384
+ - [x] Basic tests
385
+
386
+ ### Future
387
+
388
+ Potential future features include:
389
+
390
+ - `torch.compile` comparisons
391
+ - Performance regression detection
392
+ - Benchmark baselines
393
+ - CI integration
394
+ - GitHub Actions
395
+ - Additional hardware metrics
396
+ - Distributed benchmarking
397
+ - Benchmark history
398
+ - Performance dashboards
399
+
400
+ Future features are subject to change based on project feedback and development priorities.
401
+
402
+ ## Contributing
403
+
404
+ Contributions are welcome.
405
+
406
+ Before opening a pull request, please run:
407
+
408
+ ```bash
409
+ uv run pytest
410
+ uv run ruff check .
411
+ uv run ruff format --check .
412
+ ```
413
+
414
+ For larger changes, please open an issue first to discuss the proposed approach.
415
+
416
+ ## License
417
+
418
+ TorchMeter is licensed under the Apache License 2.0.
419
+
420
+ See [LICENSE](LICENSE) for the full license text.
421
+
422
+ ## Disclaimer
423
+
424
+ TorchMeter is an early-stage project. Benchmark results can vary depending on hardware, software versions, system load, and benchmark configuration.
425
+ Always compare results under controlled conditions when making performance decisions.
@@ -0,0 +1,413 @@
1
+ # TorchMeter
2
+
3
+ Simple, reproducible performance benchmarking for PyTorch models.
4
+
5
+ TorchMeter makes it easy to measure the performance of PyTorch models and generate useful benchmark results from the command line.
6
+
7
+ ## Status
8
+
9
+ 🚧 **Early Development — v0.1.0**
10
+
11
+ TorchMeter is currently an experimental open-source project.
12
+
13
+ The API and CLI may change as the project develops.
14
+
15
+ ## Features
16
+
17
+ - PyTorch model benchmarking
18
+ - CPU benchmarking
19
+ - CUDA benchmarking
20
+ - Warmup iterations
21
+ - Configurable benchmark iterations
22
+ - Mean latency
23
+ - Median latency
24
+ - P95 latency
25
+ - P99 latency
26
+ - Throughput measurement
27
+ - Peak CUDA memory measurement
28
+ - JSON output
29
+ - Markdown reports
30
+ - Simple Python API
31
+ - Reproducible environment information
32
+
33
+ ## Installation
34
+
35
+ ### Using uv
36
+
37
+ ```bash
38
+ uv add torchmeter-cli
39
+ ```
40
+
41
+ ### From source
42
+
43
+ Clone the repository:
44
+
45
+ ```bash
46
+ git clone https://github.com/YOUR_USERNAME/torchmeter-cli.git
47
+ cd torchmeter-cli
48
+ ```
49
+
50
+ Install the development environment:
51
+
52
+ ```bash
53
+ uv sync
54
+ ```
55
+
56
+ ## Quick Start
57
+
58
+ Create a Python file containing a model and example input.
59
+
60
+ For example:
61
+
62
+ ```python
63
+ import torch
64
+ import torchvision.models as models
65
+
66
+
67
+ MODEL_NAME = "ResNet50"
68
+
69
+
70
+ def create_model():
71
+ return models.resnet50(weights=None)
72
+
73
+
74
+ def create_input():
75
+ return torch.randn(1, 3, 224, 224)
76
+ ```
77
+
78
+ Save it as:
79
+
80
+ ```text
81
+ examples/model.py
82
+ ```
83
+
84
+ Run the benchmark:
85
+
86
+ ```bash
87
+ uv run torchmeter-cli run examples/model.py
88
+ ```
89
+
90
+ For CUDA:
91
+
92
+ ```bash
93
+ uv run torchmeter-cli run examples/model.py --device cuda
94
+ ```
95
+
96
+ For CPU:
97
+
98
+ ```bash
99
+ uv run torchmeter-cli run examples/model.py --device cpu
100
+ ```
101
+
102
+ ## Model Interface
103
+
104
+ TorchMeter expects the model file to provide two functions:
105
+
106
+ ```python
107
+ def create_model():
108
+ ...
109
+
110
+
111
+ def create_input():
112
+ ...
113
+ ```
114
+
115
+ You can optionally provide a model name:
116
+
117
+ ```python
118
+ MODEL_NAME = "My Model"
119
+ ```
120
+
121
+ If `MODEL_NAME` is not provided, TorchMeter uses the model class name.
122
+
123
+ ## Example Models
124
+
125
+ The repository includes several example models:
126
+
127
+ ```text
128
+ examples/
129
+ ├── cnn.py
130
+ ├── lstm.py
131
+ ├── mlp.py
132
+ ├── model.py
133
+ ├── rnn.py
134
+ ├── resnet.py
135
+ └── transformer.py
136
+ ```
137
+
138
+ Run an example:
139
+
140
+ ```bash
141
+ uv run torchmeter-cli run examples/cnn.py --device cuda
142
+ ```
143
+
144
+ ```bash
145
+ uv run torchmeter-cli run examples/rnn.py --device cuda
146
+ ```
147
+
148
+ ```bash
149
+ uv run torchmeter-cli run examples/lstm.py --device cuda
150
+ ```
151
+
152
+ ```bash
153
+ uv run torchmeter-cli run examples/transformer.py --device cuda
154
+ ```
155
+
156
+ ```bash
157
+ uv run torchmeter-cli run examples/resnet.py --device cuda
158
+ ```
159
+
160
+ ## Benchmark Configuration
161
+
162
+ Control the number of warmup iterations:
163
+
164
+ ```bash
165
+ uv run torchmeter-cli run examples/model.py --warmup 20
166
+ ```
167
+
168
+ Control the number of benchmark iterations:
169
+
170
+ ```bash
171
+ uv run torchmeter-cli run examples/model.py --iterations 100
172
+ ```
173
+
174
+ You can combine options:
175
+
176
+ ```bash
177
+ uv run torchmeter-cli run examples/model.py \
178
+ --device cuda \
179
+ --warmup 20 \
180
+ --iterations 100
181
+ ```
182
+
183
+ ## Output Formats
184
+
185
+ ### Terminal
186
+
187
+ The default output is a human-readable benchmark:
188
+
189
+ ```text
190
+ TorchMeter Benchmark
191
+ ────────────────────────────────
192
+ Model: ResNet50
193
+ Device: cuda
194
+ Iterations: 100
195
+
196
+ Latency
197
+ Mean: 4.823 ms
198
+ Median: 4.761 ms
199
+ P95: 5.132 ms
200
+ P99: 5.421 ms
201
+
202
+ Throughput: 207.34 samples/sec
203
+ Peak memory: 1420.31 MB
204
+ ```
205
+
206
+ ### JSON
207
+
208
+ Generate machine-readable output:
209
+
210
+ ```bash
211
+ uv run torchmeter-cli run examples/model.py \
212
+ --format json
213
+ ```
214
+
215
+ Write the JSON result to a file:
216
+
217
+ ```bash
218
+ uv run torchmeter-cli run examples/model.py \
219
+ --format json \
220
+ --output results.json
221
+ ```
222
+
223
+ ### Markdown
224
+
225
+ Generate a Markdown report:
226
+
227
+ ```bash
228
+ uv run torchmeter-cli run examples/model.py \
229
+ --format markdown
230
+ ```
231
+
232
+ Write it to a file:
233
+
234
+ ```bash
235
+ uv run torchmeter-cli run examples/model.py \
236
+ --format markdown \
237
+ --output benchmark.md
238
+ ```
239
+
240
+ ## Python API
241
+
242
+ TorchMeter can also be used directly from Python.
243
+
244
+ ```python
245
+ import torch
246
+
247
+ from torchmeter_cli import benchmark
248
+
249
+
250
+ model = torch.nn.Sequential(
251
+ torch.nn.Linear(768, 1024),
252
+ torch.nn.ReLU(),
253
+ torch.nn.Linear(1024, 10),
254
+ )
255
+
256
+ inputs = torch.randn(1, 768)
257
+
258
+ result = benchmark(
259
+ model,
260
+ inputs,
261
+ device="cuda",
262
+ warmup=20,
263
+ iterations=100,
264
+ )
265
+
266
+ print(result.mean_latency_ms)
267
+ print(result.throughput_samples_per_second)
268
+ ```
269
+
270
+ ## Benchmark Results
271
+
272
+ A benchmark contains:
273
+
274
+ - Model name
275
+ - Device
276
+ - Number of warmup iterations
277
+ - Number of benchmark iterations
278
+ - Mean latency
279
+ - Median latency
280
+ - P95 latency
281
+ - P99 latency
282
+ - Throughput
283
+ - Peak CUDA memory when available
284
+
285
+ ## Development
286
+
287
+ Install dependencies:
288
+
289
+ ```bash
290
+ uv sync
291
+ ```
292
+
293
+ Run tests:
294
+
295
+ ```bash
296
+ uv run pytest
297
+ ```
298
+
299
+ Run the test suite with verbose output:
300
+
301
+ ```bash
302
+ uv run pytest -v
303
+ ```
304
+
305
+ Run Ruff:
306
+
307
+ ```bash
308
+ uv run ruff check .
309
+ ```
310
+
311
+ Format the project:
312
+
313
+ ```bash
314
+ uv run ruff format .
315
+ ```
316
+
317
+ ## Project Structure
318
+
319
+ ```text
320
+ torchmeter-cli/
321
+ ├── examples/
322
+ │ ├── cnn.py
323
+ │ ├── lstm.py
324
+ │ ├── mlp.py
325
+ │ ├── model.py
326
+ │ ├── rnn.py
327
+ │ ├── resnet.py
328
+ │ └── transformer.py
329
+ │
330
+ ├── src/
331
+ │ └── torchmeter-cli/
332
+ │ ├── __init__.py
333
+ │ ├── __main__.py
334
+ │ ├── benchmark.py
335
+ │ ├── cli.py
336
+ │ └── result.py
337
+ │
338
+ ├── tests/
339
+ │ └── test_benchmark.py
340
+ │
341
+ ├── LICENSE
342
+ ├── README.md
343
+ ├── pyproject.toml
344
+ └── uv.lock
345
+ ```
346
+
347
+ ## Design Goals
348
+
349
+ TorchMeter aims to be:
350
+
351
+ 1. **Simple** — benchmark a model with a single command.
352
+ 2. **Reproducible** — capture the environment used for benchmarking.
353
+ 3. **Accurate** — correctly measure CPU and CUDA workloads.
354
+ 4. **Scriptable** — provide JSON output for automation.
355
+ 5. **Extensible** — provide both a CLI and Python API.
356
+ 6. **Open source** — developed transparently with community contributions.
357
+
358
+ ## Roadmap
359
+
360
+ ### v0.1.x
361
+
362
+ - [x] Basic PyTorch benchmarking
363
+ - [x] CPU support
364
+ - [x] CUDA support
365
+ - [x] Latency statistics
366
+ - [x] Throughput measurement
367
+ - [x] CUDA memory measurement
368
+ - [x] JSON output
369
+ - [x] Markdown output
370
+ - [x] Python API
371
+ - [x] Example models
372
+ - [x] Basic tests
373
+
374
+ ### Future
375
+
376
+ Potential future features include:
377
+
378
+ - `torch.compile` comparisons
379
+ - Performance regression detection
380
+ - Benchmark baselines
381
+ - CI integration
382
+ - GitHub Actions
383
+ - Additional hardware metrics
384
+ - Distributed benchmarking
385
+ - Benchmark history
386
+ - Performance dashboards
387
+
388
+ Future features are subject to change based on project feedback and development priorities.
389
+
390
+ ## Contributing
391
+
392
+ Contributions are welcome.
393
+
394
+ Before opening a pull request, please run:
395
+
396
+ ```bash
397
+ uv run pytest
398
+ uv run ruff check .
399
+ uv run ruff format --check .
400
+ ```
401
+
402
+ For larger changes, please open an issue first to discuss the proposed approach.
403
+
404
+ ## License
405
+
406
+ TorchMeter is licensed under the Apache License 2.0.
407
+
408
+ See [LICENSE](LICENSE) for the full license text.
409
+
410
+ ## Disclaimer
411
+
412
+ TorchMeter is an early-stage project. Benchmark results can vary depending on hardware, software versions, system load, and benchmark configuration.
413
+ Always compare results under controlled conditions when making performance decisions.
@@ -0,0 +1,30 @@
1
+ [project]
2
+ name = "torchmeter-cli"
3
+ version = "0.1.0"
4
+ description = "Simple, reproducible performance benchmarking for PyTorch models."
5
+ readme = "README.md"
6
+ requires-python = ">=3.11"
7
+ dependencies = [
8
+ "torch>=2.0",
9
+ "torchvision>=0.15",
10
+ ]
11
+
12
+ [project.license]
13
+ text = "Apache-2.0"
14
+
15
+ [[project.authors]]
16
+ name = "Pujan Neupane"
17
+ email = "pujanneupaneop0907@gmail.com"
18
+
19
+ [project.scripts]
20
+ torchmeter-cli = "torchmeter_cli.cli:main"
21
+
22
+ [build-system]
23
+ requires = ["uv_build>=0.12.1,<0.13.0"]
24
+ build-backend = "uv_build"
25
+
26
+ [dependency-groups]
27
+ dev = [
28
+ "pytest>=9.1.1",
29
+ "ruff>=0.16.2",
30
+ ]
@@ -0,0 +1,20 @@
1
+ [project]
2
+ name = "torchmeter-cli"
3
+ version = "0.1.0"
4
+ description = "Simple, reproducible performance benchmarking for PyTorch models."
5
+ readme = "README.md"
6
+ license = { text = "Apache-2.0" }
7
+ authors = [{ name = "Pujan Neupane", email = "pujanneupaneop0907@gmail.com" }]
8
+ requires-python = ">=3.11"
9
+ dependencies = ["torch>=2.0", "torchvision>=0.15"]
10
+
11
+ [project.scripts]
12
+ torchmeter-cli = "torchmeter_cli.cli:main"
13
+
14
+ [build-system]
15
+ requires = ["uv_build>=0.12.1,<0.13.0"]
16
+
17
+ build-backend = "uv_build"
18
+
19
+ [dependency-groups]
20
+ dev = ["pytest>=9.1.1", "ruff>=0.16.2"]
@@ -0,0 +1,11 @@
1
+ """TorchMeter CLI: PyTorch model performance benchmarking."""
2
+
3
+ from .benchmark import benchmark
4
+ from .result import BenchmarkResult
5
+
6
+ __version__ = "0.1.0"
7
+
8
+ __all__ = [
9
+ "BenchmarkResult",
10
+ "benchmark",
11
+ ]
@@ -0,0 +1,4 @@
1
+ from .cli import main
2
+
3
+ if __name__ == "__main__":
4
+ main()
@@ -0,0 +1,196 @@
1
+ from __future__ import annotations
2
+
3
+ import statistics
4
+ import time
5
+ from collections.abc import Callable
6
+
7
+ import torch
8
+
9
+ from .result import BenchmarkResult
10
+
11
+
12
+ def _synchronize(device: torch.device) -> None:
13
+ """Synchronize CUDA so GPU operations are fully completed."""
14
+ if device.type == "cuda":
15
+ torch.cuda.synchronize(device)
16
+
17
+
18
+ def _percentile(values: list[float], percentile: float) -> float:
19
+ """Calculate a percentile using linear interpolation."""
20
+ if not values:
21
+ raise ValueError("Cannot calculate percentile from empty data.")
22
+
23
+ values = sorted(values)
24
+
25
+ if len(values) == 1:
26
+ return values[0]
27
+
28
+ index = (len(values) - 1) * percentile
29
+ lower = int(index)
30
+ upper = min(lower + 1, len(values) - 1)
31
+
32
+ weight = index - lower
33
+
34
+ return values[lower] + (values[upper] - values[lower]) * weight
35
+
36
+ def _set_batch_size(inputs: object, batch_size: int) -> object:
37
+ """Create inputs with the requested batch size."""
38
+ if torch.is_tensor(inputs):
39
+ if inputs.ndim == 0:
40
+ return inputs
41
+
42
+ shape = list(inputs.shape)
43
+ shape[0] = batch_size
44
+
45
+ return torch.randn(
46
+ shape,
47
+ dtype=inputs.dtype,
48
+ device=inputs.device,
49
+ )
50
+
51
+ if isinstance(inputs, dict):
52
+ return {
53
+ key: _set_batch_size(value, batch_size)
54
+ for key, value in inputs.items()
55
+ }
56
+
57
+ if isinstance(inputs, tuple):
58
+ return tuple(
59
+ _set_batch_size(value, batch_size)
60
+ for value in inputs
61
+ )
62
+
63
+ if isinstance(inputs, list):
64
+ return [
65
+ _set_batch_size(value, batch_size)
66
+ for value in inputs
67
+ ]
68
+
69
+ return inputs
70
+ def _run_model(
71
+ model: torch.nn.Module,
72
+ inputs: object,
73
+ ) -> object:
74
+ """Run the model with the provided inputs."""
75
+ if isinstance(inputs, dict):
76
+ return model(**inputs)
77
+
78
+ if isinstance(inputs, (tuple, list)):
79
+ return model(*inputs)
80
+
81
+ return model(inputs)
82
+
83
+
84
+ def _move_to_device(
85
+ value: object,
86
+ device: torch.device,
87
+ ) -> object:
88
+ """Move tensors or nested structures to the target device."""
89
+ if torch.is_tensor(value):
90
+ return value.to(device)
91
+
92
+ if isinstance(value, dict):
93
+ return {
94
+ key: _move_to_device(item, device)
95
+ for key, item in value.items()
96
+ }
97
+
98
+ if isinstance(value, tuple):
99
+ return tuple(_move_to_device(item, device) for item in value)
100
+
101
+ if isinstance(value, list):
102
+ return [_move_to_device(item, device) for item in value]
103
+
104
+ return value
105
+
106
+
107
+ def benchmark(
108
+ model: torch.nn.Module,
109
+ inputs: object,
110
+ *,
111
+ model_name: str = "PyTorch Model",
112
+ device: str = "auto",
113
+ warmup: int = 20,
114
+ iterations: int = 100,
115
+ ) -> BenchmarkResult:
116
+ """Benchmark a PyTorch model.
117
+
118
+ Args:
119
+ model: PyTorch model to benchmark.
120
+ inputs: Example model input.
121
+ model_name: Human-readable model name.
122
+ device: "auto", "cpu", or "cuda".
123
+ warmup: Number of warmup iterations.
124
+ iterations: Number of measured iterations.
125
+ """
126
+ if warmup < 0:
127
+ raise ValueError("warmup must be >= 0")
128
+
129
+ if iterations <= 0:
130
+ raise ValueError("iterations must be > 0")
131
+
132
+ if device == "auto":
133
+ device = "cuda" if torch.cuda.is_available() else "cpu"
134
+
135
+ target_device = torch.device(device)
136
+
137
+ if target_device.type == "cuda" and not torch.cuda.is_available():
138
+ raise RuntimeError("CUDA was requested but is not available.")
139
+
140
+ model = model.to(target_device)
141
+ model.eval()
142
+
143
+ inputs = _move_to_device(inputs, target_device)
144
+
145
+ if target_device.type == "cuda":
146
+ torch.cuda.reset_peak_memory_stats(target_device)
147
+
148
+ with torch.inference_mode():
149
+ for _ in range(warmup):
150
+ _run_model(model, inputs)
151
+
152
+ _synchronize(target_device)
153
+
154
+ timings: list[float] = []
155
+
156
+ for _ in range(iterations):
157
+ _synchronize(target_device)
158
+
159
+ start = time.perf_counter()
160
+
161
+ _run_model(model, inputs)
162
+
163
+ _synchronize(target_device)
164
+
165
+ end = time.perf_counter()
166
+
167
+ timings.append((end - start) * 1000)
168
+
169
+ mean_latency = statistics.mean(timings)
170
+ median_latency = statistics.median(timings)
171
+
172
+ p95 = _percentile(timings, 0.95)
173
+ p99 = _percentile(timings, 0.99)
174
+
175
+ throughput = 1000.0 / mean_latency
176
+
177
+ peak_memory = None
178
+
179
+ if target_device.type == "cuda":
180
+ peak_memory = (
181
+ torch.cuda.max_memory_allocated(target_device)
182
+ / (1024**2)
183
+ )
184
+
185
+ return BenchmarkResult(
186
+ model_name=model_name,
187
+ device=str(target_device),
188
+ iterations=iterations,
189
+ warmup_iterations=warmup,
190
+ mean_latency_ms=mean_latency,
191
+ median_latency_ms=median_latency,
192
+ p95_latency_ms=p95,
193
+ p99_latency_ms=p99,
194
+ throughput_samples_per_second=throughput,
195
+ peak_memory_mb=peak_memory,
196
+ )
@@ -0,0 +1,169 @@
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import importlib.util
5
+ import sys
6
+ from pathlib import Path
7
+
8
+ from .benchmark import benchmark
9
+
10
+
11
+ def _load_module(path: Path):
12
+ """Load a Python file as a module."""
13
+ spec = importlib.util.spec_from_file_location(
14
+ "torchmeter_cli_user_model",
15
+ path,
16
+ )
17
+
18
+ if spec is None or spec.loader is None:
19
+ raise RuntimeError(f"Could not load {path}")
20
+
21
+ module = importlib.util.module_from_spec(spec)
22
+
23
+ sys.path.insert(0, str(path.parent))
24
+
25
+ try:
26
+ spec.loader.exec_module(module)
27
+ finally:
28
+ sys.path.pop(0)
29
+
30
+ return module
31
+
32
+
33
+ def _build_parser() -> argparse.ArgumentParser:
34
+ parser = argparse.ArgumentParser(
35
+ prog="torchmeter-cli",
36
+ description="Performance benchmarking for PyTorch models.",
37
+ )
38
+
39
+ subparsers = parser.add_subparsers(dest="command", required=True)
40
+
41
+ run_parser = subparsers.add_parser(
42
+ "run",
43
+ help="Benchmark a Python model file.",
44
+ )
45
+
46
+ run_parser.add_argument(
47
+ "model",
48
+ type=Path,
49
+ help="Path to the model.py file.",
50
+ )
51
+
52
+ run_parser.add_argument(
53
+ "--device",
54
+ choices=["auto", "cpu", "cuda"],
55
+ default="auto",
56
+ )
57
+
58
+ run_parser.add_argument(
59
+ "--warmup",
60
+ type=int,
61
+ default=20,
62
+ )
63
+
64
+ run_parser.add_argument(
65
+ "--iterations",
66
+ type=int,
67
+ default=100,
68
+ )
69
+
70
+ run_parser.add_argument(
71
+ "--format",
72
+ choices=["text", "json", "markdown"],
73
+ default="text",
74
+ )
75
+
76
+ run_parser.add_argument(
77
+ "--output",
78
+ type=Path,
79
+ help="Write the result to a file.",
80
+ )
81
+ run_parser.add_argument(
82
+ "--batch-sizes",
83
+ type=str,
84
+ help="Comma-separated batch sizes to benchmark, e.g. 1,8,16,32.",
85
+ )
86
+ return parser
87
+
88
+
89
+ def _print_text(result) -> None:
90
+ print()
91
+ print("TorchMeter Benchmark")
92
+ print("─" * 32)
93
+ print(f"Model: {result.model_name}")
94
+ print(f"Device: {result.device}")
95
+ print(f"Iterations: {result.iterations}")
96
+ print()
97
+ print("Latency")
98
+ print(f" Mean: {result.mean_latency_ms:.3f} ms")
99
+ print(f" Median: {result.median_latency_ms:.3f} ms")
100
+ print(f" P95: {result.p95_latency_ms:.3f} ms")
101
+ print(f" P99: {result.p99_latency_ms:.3f} ms")
102
+ print()
103
+ print(f"Throughput: {result.throughput_samples_per_second:.2f} samples/sec")
104
+
105
+ if result.peak_memory_mb is not None:
106
+ print(f"Peak memory: {result.peak_memory_mb:.2f} MB")
107
+
108
+ print()
109
+
110
+
111
+ def main() -> None:
112
+ parser = _build_parser()
113
+ args = parser.parse_args()
114
+
115
+ if args.command != "run":
116
+ parser.print_help()
117
+ return
118
+
119
+ if not args.model.exists():
120
+ parser.error(f"Model file does not exist: {args.model}")
121
+
122
+ try:
123
+ module = _load_module(args.model)
124
+
125
+ if not hasattr(module, "create_model"):
126
+ parser.error("model.py must define create_model().")
127
+
128
+ if not hasattr(module, "create_input"):
129
+ parser.error("model.py must define create_input().")
130
+
131
+ model = module.create_model()
132
+ inputs = module.create_input()
133
+
134
+ model_name = getattr(
135
+ module,
136
+ "MODEL_NAME",
137
+ model.__class__.__name__,
138
+ )
139
+
140
+ result = benchmark(
141
+ model,
142
+ inputs,
143
+ model_name=model_name,
144
+ device=args.device,
145
+ warmup=args.warmup,
146
+ iterations=args.iterations,
147
+ )
148
+
149
+ except Exception as exc:
150
+ parser.error(str(exc))
151
+
152
+ if args.format == "json":
153
+ output = result.to_json()
154
+ elif args.format == "markdown":
155
+ output = result.to_markdown()
156
+ else:
157
+ _print_text(result)
158
+ output = None
159
+
160
+ if output is not None:
161
+ if args.output:
162
+ args.output.write_text(output + "\n")
163
+ print(f"Report written to {args.output}")
164
+ else:
165
+ print(output)
166
+
167
+
168
+ if __name__ == "__main__":
169
+ main()
@@ -0,0 +1,47 @@
1
+ from dataclasses import asdict, dataclass
2
+ import json
3
+
4
+
5
+ @dataclass
6
+ class BenchmarkResult:
7
+ model_name: str
8
+ device: str
9
+ iterations: int
10
+ warmup_iterations: int
11
+
12
+ mean_latency_ms: float
13
+ median_latency_ms: float
14
+ p95_latency_ms: float
15
+ p99_latency_ms: float
16
+
17
+ throughput_samples_per_second: float
18
+ peak_memory_mb: float | None = None
19
+
20
+ def to_dict(self) -> dict:
21
+ return asdict(self)
22
+
23
+ def to_json(self) -> str:
24
+ return json.dumps(self.to_dict(), indent=2)
25
+
26
+ def to_markdown(self) -> str:
27
+ memory = (
28
+ f"{self.peak_memory_mb:.2f} MB"
29
+ if self.peak_memory_mb is not None
30
+ else "N/A"
31
+ )
32
+
33
+ return f"""# TorchMeter Benchmark
34
+
35
+ | Metric | Value |
36
+ |---|---:|
37
+ | Model | `{self.model_name}` |
38
+ | Device | `{self.device}` |
39
+ | Iterations | {self.iterations} |
40
+ | Warmup | {self.warmup_iterations} |
41
+ | Mean latency | {self.mean_latency_ms:.3f} ms |
42
+ | Median latency | {self.median_latency_ms:.3f} ms |
43
+ | P95 latency | {self.p95_latency_ms:.3f} ms |
44
+ | P99 latency | {self.p99_latency_ms:.3f} ms |
45
+ | Throughput | {self.throughput_samples_per_second:.2f} samples/sec |
46
+ | Peak GPU memory | {memory} |
47
+ """