|
|
|
|
|
|
|
|
|
|
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python import core, workspace |
| from caffe2.python.models.download import ModelDownloader |
| import numpy as np |
| import argparse |
| import time |
|
|
|
|
| def GetArgumentParser(): |
| parser = argparse.ArgumentParser(description="Caffe2 benchmark.") |
| parser.add_argument( |
| "--batch_size", |
| type=int, |
| default=128, |
| help="The batch size." |
| ) |
| parser.add_argument("--model", type=str, help="The model to benchmark.") |
| parser.add_argument( |
| "--order", |
| type=str, |
| default="NCHW", |
| help="The order to evaluate." |
| ) |
| parser.add_argument( |
| "--device", |
| type=str, |
| default="CPU", |
| help="device to evaluate on." |
| ) |
| parser.add_argument( |
| "--cudnn_ws", |
| type=int, |
| help="The cudnn workspace size." |
| ) |
| parser.add_argument( |
| "--iterations", |
| type=int, |
| default=10, |
| help="Number of iterations to run the network." |
| ) |
| parser.add_argument( |
| "--warmup_iterations", |
| type=int, |
| default=10, |
| help="Number of warm-up iterations before benchmarking." |
| ) |
| parser.add_argument( |
| "--forward_only", |
| action='store_true', |
| help="If set, only run the forward pass." |
| ) |
| parser.add_argument( |
| "--layer_wise_benchmark", |
| action='store_true', |
| help="If True, run the layer-wise benchmark as well." |
| ) |
| parser.add_argument( |
| "--engine", |
| type=str, |
| default="", |
| help="If set, blindly prefer the given engine(s) for every op.") |
| parser.add_argument( |
| "--dump_model", |
| action='store_true', |
| help="If True, dump the model prototxts to disk." |
| ) |
| parser.add_argument("--net_type", type=str, default="simple") |
| parser.add_argument("--num_workers", type=int, default=2) |
| parser.add_argument("--use-nvtx", default=False, action='store_true') |
| parser.add_argument("--htrace_span_log_path", type=str) |
| return parser |
|
|
|
|
| def benchmark(args): |
| print('Batch size: {}'.format(args.batch_size)) |
| mf = ModelDownloader() |
| init_net, pred_net, value_info = mf.get_c2_model(args.model) |
| input_shapes = {k : [args.batch_size] + v[-1][1:] for (k, v) in value_info.items()} |
| print("input info: {}".format(input_shapes)) |
| external_inputs = {} |
| for k, v in input_shapes.items(): |
| external_inputs[k] = np.random.randn(*v).astype(np.float32) |
|
|
| if args.device == 'CPU': |
| device_option = core.DeviceOption(caffe2_pb2.CPU) |
| elif args.device == 'MKL': |
| device_option = core.DeviceOption(caffe2_pb2.MKLDNN) |
| elif args.device == 'IDEEP': |
| device_option = core.DeviceOption(caffe2_pb2.IDEEP) |
| else: |
| raise Exception("Unknown device: {}".format(args.device)) |
| print("Device option: {}, {}".format(args.device, device_option)) |
| pred_net.device_option.CopyFrom(device_option) |
| for op in pred_net.op: |
| op.device_option.CopyFrom(device_option) |
|
|
| |
| workspace.RunNetOnce(init_net) |
| bb = workspace.Blobs() |
| weights = {} |
| for b in bb: |
| weights[b] = workspace.FetchBlob(b) |
| for k, v in external_inputs.items(): |
| weights[k] = v |
| workspace.ResetWorkspace() |
|
|
| with core.DeviceScope(device_option): |
| for name, blob in weights.items(): |
| |
| workspace.FeedBlob(name, blob, device_option) |
| workspace.CreateNet(pred_net) |
| start = time.time() |
| res = workspace.BenchmarkNet(pred_net.name, |
| args.warmup_iterations, |
| args.iterations, |
| args.layer_wise_benchmark) |
| print("FPS: {:.2f}".format(1/res[0]*1000*args.batch_size)) |
|
|
| if __name__ == '__main__': |
| args, extra_args = GetArgumentParser().parse_known_args() |
| if ( |
| not args.batch_size or not args.model or not args.order |
| ): |
| GetArgumentParser().print_help() |
| benchmark(args) |
|
|
