mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
a229b4526f
Preferring dash over underscore in command-line options. Add `--command-arg-name` to the argument parser. The old arguments with underscores `--command_arg_name` are kept for backward compatibility.
Both dashes and underscores are used in the PyTorch codebase. Some argument parsers only have dashes or only have underscores in arguments. For example, the `torchrun` utility for distributed training only accepts underscore arguments (e.g., `--master_port`). The dashes are more common in other command-line tools. And it looks to be the default choice in the Python standard library:
`argparse.BooleanOptionalAction`: 4a9dff0e5a/Lib/argparse.py (L893-L895)
```python
class BooleanOptionalAction(Action):
def __init__(...):
if option_string.startswith('--'):
option_string = '--no-' + option_string[2:]
_option_strings.append(option_string)
```
It adds `--no-argname`, not `--no_argname`. Also typing `_` need to press the shift or the caps-lock key than `-`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/94505
Approved by: https://github.com/ezyang, https://github.com/seemethere
104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
import argparse
|
|
import sys
|
|
import torch
|
|
from .utils import gen_sparse_csr, gen_sparse_coo, gen_sparse_coo_and_csr, Event
|
|
|
|
def test_sparse_csr(m, nnz, test_count):
|
|
start_timer = Event(enable_timing=True)
|
|
stop_timer = Event(enable_timing=True)
|
|
|
|
csr = gen_sparse_csr((m, m), nnz)
|
|
vector = torch.randn(m, dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start_timer.record()
|
|
csr.matmul(vector)
|
|
stop_timer.record()
|
|
times.append(start_timer.elapsed_time(stop_timer))
|
|
|
|
return sum(times) / len(times)
|
|
|
|
def test_sparse_coo(m, nnz, test_count):
|
|
start_timer = Event(enable_timing=True)
|
|
stop_timer = Event(enable_timing=True)
|
|
|
|
coo = gen_sparse_coo((m, m), nnz)
|
|
vector = torch.randn(m, dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start_timer.record()
|
|
coo.matmul(vector)
|
|
stop_timer.record()
|
|
times.append(start_timer.elapsed_time(stop_timer))
|
|
|
|
return sum(times) / len(times)
|
|
|
|
def test_sparse_coo_and_csr(m, nnz, test_count):
|
|
start = Event(enable_timing=True)
|
|
stop = Event(enable_timing=True)
|
|
|
|
coo, csr = gen_sparse_coo_and_csr((m, m), nnz)
|
|
vector = torch.randn(m, dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start.record()
|
|
coo.matmul(vector)
|
|
stop.record()
|
|
|
|
times.append(start.elapsed_time(stop))
|
|
|
|
coo_mean_time = sum(times) / len(times)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start.record()
|
|
csr.matmul(vector)
|
|
stop.record()
|
|
times.append(start.elapsed_time(stop))
|
|
|
|
csr_mean_time = sum(times) / len(times)
|
|
|
|
return coo_mean_time, csr_mean_time
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="SpMV")
|
|
|
|
parser.add_argument("--format", default='csr', type=str)
|
|
parser.add_argument("--m", default='1000', type=int)
|
|
parser.add_argument("--nnz-ratio", "--nnz_ratio", default='0.1', type=float)
|
|
parser.add_argument("--outfile", default='stdout', type=str)
|
|
parser.add_argument("--test-count", "--test_count", default='10', type=int)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.outfile == 'stdout':
|
|
outfile = sys.stdout
|
|
elif args.outfile == 'stderr':
|
|
outfile = sys.stderr
|
|
else:
|
|
outfile = open(args.outfile, "a")
|
|
|
|
test_count = args.test_count
|
|
m = args.m
|
|
nnz_ratio = args.nnz_ratio
|
|
|
|
nnz = int(nnz_ratio * m * m)
|
|
if args.format == 'csr':
|
|
time = test_sparse_csr(m, nnz, test_count)
|
|
elif args.format == 'coo':
|
|
time = test_sparse_coo(m, nnz, test_count)
|
|
elif args.format == 'both':
|
|
time_coo, time_csr = test_sparse_coo_and_csr(m, nnz, test_count)
|
|
|
|
if args.format != 'both':
|
|
print("format=", args.format, " nnz_ratio=", nnz_ratio, " m=", m,
|
|
" time=", time, file=outfile)
|
|
else:
|
|
print("format=coo", " nnz_ratio=", nnz_ratio, " m=", m,
|
|
" time=", time_coo, file=outfile)
|
|
print("format=csr", " nnz_ratio=", nnz_ratio, " m=", m,
|
|
" time=", time_csr, file=outfile)
|