bench.py

import argparse
import copy
import gc
import json
import sys
import time
from collections import namedtuple

import torch
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 f"{item:.4g}"
    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:
    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(f"Benchmarking {bench_name}s...")
    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",
        "--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",
        "--cuda_pointwise_loop_level",
        default=None,
        type=int,
    )
    parser.add_argument(
        "--cuda-pointwise-block-count",
        "--cuda_pointwise_block_count",
        default=None,
        type=int,
    )
    parser.add_argument(
        "--cuda-pointwise-block-size",
        "--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)