import argparse from collections import namedtuple import torch import gc import sys import json import copy import time from torch.autograd.profiler import record_function from .fuser import set_fuser from .runner import get_nn_runners BenchResult = namedtuple('BenchResult', [ 'name', 'avg_fwd', 'std_fwd', 'info_fwd', 'avg_bwd', 'std_bwd', 'info_bwd', ]) def fit_str(string, colwidth=16): if len(string) < colwidth: return (colwidth - len(string)) * ' ' + string else: return string[:colwidth] def to_str(item): if isinstance(item, float): return '%.4g' % item return str(item) def print_header(colwidth=16, sep=' '): items = [] for item in BenchResult._fields: items.append(fit_str(item)) return sep.join(items) def pretty_print(benchresult, colwidth=16, sep=' '): items = [] for thing in benchresult: items.append(fit_str(to_str(thing))) return sep.join(items) # shim for torch.cuda.Event when running on cpu class Event(object): def __init__(self, enable_timing): pass def record(self): self.time = time.perf_counter() def elapsed_time(self, end_event): assert isinstance(end_event, Event) return end_event.time - self.time def trainbench(name, rnn_creator, nloops=100, warmup=10, seqLength=100, numLayers=1, inputSize=512, hiddenSize=512, miniBatch=64, device='cuda', seed=None): def train_batch(modeldef): # CUDA events for timing if device == 'cuda': timer_class = torch.cuda.Event else: timer_class = Event fwd_start_event = timer_class(enable_timing=True) fwd_end_event = timer_class(enable_timing=True) bwd_start_event = timer_class(enable_timing=True) bwd_end_event = timer_class(enable_timing=True) gc.collect() fwd_start_event.record() with record_function("## forward ##"): forward_output = modeldef.forward(*modeldef.inputs) fwd_end_event.record() # XXX: Use if need to print something # print(modeldef.forward.graph_for(*modeldef.inputs)) if modeldef.backward_setup is not None: backward_input = modeldef.backward_setup(forward_output) else: backward_input = forward_output gc.collect() bwd_start_event.record() if modeldef.backward is not None: modeldef.backward(*backward_input) bwd_end_event.record() if modeldef.backward is not None: with torch.no_grad(): for param in modeldef.params: assert param.grad is not None param.grad.zero_() if device == 'cuda': torch.cuda.synchronize() fwd_time = fwd_start_event.elapsed_time(fwd_end_event) bwd_time = bwd_start_event.elapsed_time(bwd_end_event) return fwd_time, bwd_time creator_args = creator_args = { 'seqLength': seqLength, 'numLayers': numLayers, 'inputSize': inputSize, 'hiddenSize': hiddenSize, 'miniBatch': miniBatch, 'device': device, 'seed': seed } modeldef = rnn_creator(**creator_args) [train_batch(modeldef) for _ in range(warmup)] results = [train_batch(modeldef) for _ in range(nloops)] fwd_times, bwd_times = zip(*results) fwd_times = torch.tensor(fwd_times) bwd_times = torch.tensor(bwd_times) return BenchResult(name=name, avg_fwd=fwd_times.mean().item(), std_fwd=fwd_times.std().item(), info_fwd=fwd_times, avg_bwd=bwd_times.mean().item(), std_bwd=bwd_times.std().item(), info_bwd=bwd_times) def print_stderr(*args, **kwargs): kwargs['file'] = sys.stderr return print(*args, **kwargs) def print_json_oss_format(results): oss_results = {} for group_name, group_val in results.items(): oss_results[group_name] = {} for model_name, run_time in group_val.items(): # Output for OSS oss_results[group_name][model_name] = run_time['avg'] print(json.dumps(oss_results)) def print_json_pep_format(results): # print the AI-PEP format json string for each model for group_name, group_val in results.items(): for model_name, run_time in group_val.items(): # Output for AI-PEP num_iters = len(run_time['info']) info = run_time['info'].tolist() for i in range(num_iters): print("Caffe2Observer " + json.dumps( { "type": "NET", "metric": group_name + "-" + model_name, "unit": "ms", "value": str(info[i]) } )) def bench(rnn_runners, group_name, print_json=False, sep=' ', **params): print_stderr(print_header(sep=sep)) results = {} for name, creator, context in rnn_runners: with context(): try: result = trainbench(name, creator, **params) # Replace the value of info_fwd and info_bwd to None result_with_no_info = result._replace( info_fwd='None', info_bwd='None') print_stderr(pretty_print(result_with_no_info, sep=sep)) results[name] = result except Exception as e: if not print_json: raise return { group_name: {k: {"avg": v.avg_fwd, "std": v.std_fwd, "info": v.info_fwd} for k, v in results.items()}, group_name + '-backward': {k: {"avg": v.avg_bwd, "std": v.std_bwd, "info": v.info_bwd} for k, v in results.items()}, } def bench_group(model_list, bench_name, bench_group, bench_args): print_stderr('Benchmarking {}s...'.format(bench_name)) nn_results = bench(get_nn_runners(*model_list), bench_group, **bench_args) print_stderr('') return nn_results if __name__ == '__main__': parser = argparse.ArgumentParser(description='Profile RNNs') # groups help control which test group you want to run # if you only want to run one/two benchmark, run it with # e.g: python -m fastrnns.bench --rnns jit and --group rnns default_groups = ['cnns', 'rnns'] parser.add_argument('--seqLength', default='100', type=int) parser.add_argument('--numLayers', default='1', type=int) parser.add_argument('--inputSize', default='512', type=int) parser.add_argument('--hiddenSize', default='512', type=int) parser.add_argument('--miniBatch', default='64', type=int) parser.add_argument('--warmup', default='10', type=int) parser.add_argument('--nloops', default='100', type=int) parser.add_argument('--device', default='cuda', type=str) parser.add_argument('--variable_lstms', action='store_true', help='Also benchmark variable sequence length lstms ' 'Note that some of these run really slowly ' 'and that the `seqLength` flag will be ignored.') parser.add_argument('--sep', default=' ', type=str) parser.add_argument('--print-json', nargs='?', default=None, const='oss') parser.add_argument('--rnns', nargs='*', help='What to run. cudnn, aten, jit, etc') parser.add_argument('--cnns', nargs='*', help='What to run. resnet18, resnet18_jit, resnet50, etc') parser.add_argument('--group', nargs='*', default=default_groups, help='Which group to run. cnns, rnns, etc.') parser.add_argument('--fuser', default='te', type=str, help='The fuser backend to use. One of: te, old, or none') parser.add_argument('--executor', default=None, type=str, help='The executor to use. One of: legacy, simple, profiling') parser.add_argument('--cuda_pointwise_loop_level', default=None, type=int) parser.add_argument('--cuda_pointwise_block_count', default=None, type=int) parser.add_argument('--cuda_pointwise_block_size', default=None, type=int) args = parser.parse_args() set_fuser(args.fuser, args.executor) if args.cuda_pointwise_loop_level: torch._C._jit_set_te_cuda_pointwise_loop_levels(args.cuda_pointwise_loop_level) if args.cuda_pointwise_block_count: torch._C._jit_set_te_cuda_pointwise_block_count(args.cuda_pointwise_block_count) if args.cuda_pointwise_block_size: torch._C._jit_set_te_cuda_pointwise_block_size(args.cuda_pointwise_block_size) rnns = args.rnns or ['cudnn', 'aten', 'jit', 'jit_premul', 'jit_premul_bias', 'jit_simple', 'jit_multilayer', 'py'] cnns = args.cnns or ['resnet18', 'resnet18_jit', 'resnet50', 'resnet50_jit'] # TODO: Maybe add a separate section for the layernorm/dropout lstms # 'cudnn_layernorm', jit_layernorm', 'jit_layernom_decom', # 'jit', 'jit_dropout', 'cudnn_dropout' vlrnns = ['vl_cudnn', 'vl_jit', 'vl_py'] if args.print_json: print_stderr = lambda *args, **kwargs: None # noqa: E731,F811 print_stderr(args) bench_args = copy.deepcopy(vars(args)) should_bench_varlen_lstms = args.variable_lstms del bench_args['group'] del bench_args['rnns'] del bench_args['cnns'] del bench_args['variable_lstms'] del bench_args['fuser'] del bench_args['executor'] del bench_args['cuda_pointwise_loop_level'] del bench_args['cuda_pointwise_block_count'] del bench_args['cuda_pointwise_block_size'] results = {} if should_bench_varlen_lstms: if args.nloops + args.warmup > 30: print_stderr( 'WARNING: some of the variable sequence length lstms are ' 'very unoptimized and therefore take forever to run.') results.update(bench_group(vlrnns, 'variable-length sequence LSTM', 'vl_lstm', bench_args)) if 'rnns' in args.group: results.update(bench_group(rnns, 'LSTM', 'lstm', bench_args)) if 'cnns' in args.group: results.update(bench_group(cnns, 'ResNet', 'resnet', bench_args)) if args.print_json == 'oss': print_json_oss_format(results) elif args.print_json == 'pep': print_json_pep_format(results)