mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +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
106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
import argparse
|
|
import sys
|
|
import torch
|
|
from utils import gen_sparse_csr, gen_sparse_coo, Event
|
|
|
|
def test_sparse_csr(m, n, k, nnz, test_count):
|
|
start_timer = Event(enable_timing=True)
|
|
stop_timer = Event(enable_timing=True)
|
|
|
|
csr = gen_sparse_csr((m, k), nnz)
|
|
mat = torch.randn(k, n, dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start_timer.record()
|
|
csr.matmul(mat)
|
|
stop_timer.record()
|
|
times.append(start_timer.elapsed_time(stop_timer))
|
|
|
|
return sum(times) / len(times)
|
|
|
|
def test_sparse_coo(m, n, k, nnz, test_count):
|
|
start_timer = Event(enable_timing=True)
|
|
stop_timer = Event(enable_timing=True)
|
|
|
|
coo = gen_sparse_coo((m, k), nnz)
|
|
mat = torch.randn(k, n, dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start_timer.record()
|
|
coo.matmul(mat)
|
|
stop_timer.record()
|
|
times.append(start_timer.elapsed_time(stop_timer))
|
|
|
|
return sum(times) / len(times)
|
|
|
|
def test_sparse_coo_and_csr(m, n, k, nnz, test_count):
|
|
start = Event(enable_timing=True)
|
|
stop = Event(enable_timing=True)
|
|
|
|
coo, csr = gen_sparse_coo_and_csr((m, k), nnz)
|
|
mat = torch.randn((k, n), dtype=torch.double)
|
|
|
|
times = []
|
|
for _ in range(test_count):
|
|
start.record()
|
|
coo.matmul(mat)
|
|
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(mat)
|
|
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="SpMM")
|
|
|
|
parser.add_argument("--format", default='csr', type=str)
|
|
parser.add_argument("--m", default='1000', type=int)
|
|
parser.add_argument("--n", default='1000', type=int)
|
|
parser.add_argument("--k", 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
|
|
n = args.n
|
|
k = args.k
|
|
nnz_ratio = args.nnz_ratio
|
|
|
|
nnz = int(nnz_ratio * m * k)
|
|
if args.format == 'csr':
|
|
time = test_sparse_csr(m, n, k, nnz, test_count)
|
|
elif args.format == 'coo':
|
|
time = test_sparse_coo(m, n, k, 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=coo", " nnz_ratio=", nnz_ratio, " m=", m, " n=", n, " k=", k, " time=", time_coo, file=outfile)
|
|
print("format=csr", " nnz_ratio=", nnz_ratio, " m=", m, " n=", n, " k=", k, " time=", time_csr, file=outfile)
|
|
else:
|
|
print("format=", args.format, " nnz_ratio=", nnz_ratio, " m=", m, " n=", n, " k=", k, " time=", time,
|
|
file=outfile)
|