1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
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)
# Hack to initialized weights into MKL/IDEEP context
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():
#print("{}".format(name))
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)
|
