Caffe2 - Python API
A deep learning, cross platform ML framework
benchmark_generator.py
1 # Copyright (c) 2016-present, Facebook, Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 ##############################################################################
15 
16 #!/usr/bin/env python
17 
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 from __future__ import unicode_literals
22 import string
23 
24 import argparse
25 
26 import numpy as np
27 
28 from caffe2.python.model_helper import ModelHelper
29 from caffe2.python.predictor import mobile_exporter
30 from caffe2.python import core, workspace, brew, utils
31 
32 
33 def parse_kwarg(kwarg_str):
34  key, value = map(string.strip, kwarg_str.split("=", 1))
35  try:
36  value = int(value)
37  except ValueError:
38  try:
39  value = float(value)
40  except ValueError:
41  pass
42  return key, value
43 
44 
45 def main(args):
46  # User defined keyword arguments
47  kwargs = {"order": "NCHW"}
48  kwargs.update(dict(args.kwargs))
49 
50  model = ModelHelper(name=args.benchmark_name)
51 
52  op_type = args.operator # assumes a brew type op name
53  input_name = args.input_name
54  output_name = args.output_name
55 
56  iters = int(args.iters)
57  for i in range(iters):
58  input_blob_name = input_name + (str(i) if i > 0 and args.chain else '')
59  output_blob_name = output_name + str(i + 1)
60  add_op = getattr(brew, op_type)
61  add_op(model, input_blob_name, output_blob_name, **kwargs)
62  if args.chain:
63  input_name, output_name = output_name, input_name
64 
65  workspace.RunNetOnce(model.param_init_net)
66  extra_init_net_ops = []
67 
68  def make_blob_on_context(blob_name, blob_data, context):
69  if context.upper() != "CPU":
70  blob_name_modified = "{}_CPU".format(blob_name)
71  else: # CPU case is simple
72  blob_name_modified = blob_name
73 
74  fill_op = core.CreateOperator(
75  "GivenTensorFill", [], [blob_name_modified],
76  arg=[
77  utils.MakeArgument("shape", blob_data.shape),
78  utils.MakeArgument("values", blob_data)
79  ]
80  )
81  extra_init_net_ops.append(fill_op)
82 
83  # We need to create CPU blobs and add some copy operations in
84  # the init_net
85  if context.upper() == "OPENGL":
86  copy_op = core.CreateOperator("CopyToOpenGL", [blob_name_modified],
87  [blob_name])
88  extra_init_net_ops.append(copy_op)
89 
90  for unparsed_blob in args.blob:
91  name, unparsed_dims = unparsed_blob.split('=')
92  dims = [int(d) for d in unparsed_dims.split(',')]
93  np_input = np.random.rand(*dims).astype(np.float32)
94  make_blob_on_context(name, np_input, args.context)
95 
96  init_net, predict_net = mobile_exporter.Export(
97  workspace, model.net, model.params
98  )
99  init_net.op.extend(extra_init_net_ops)
100 
101  # Handle manual rewrite
102  if args.context.upper() == "OPENGL":
103  old_ops = [op for op in predict_net.op]
104  del predict_net.op[:]
105  for op in old_ops:
106  op.type = 'OpenGL{}'.format(op.type)
107  predict_net.op.extend(old_ops)
108 
109  if args.debug:
110  print("init_net:")
111  for op in init_net.op:
112  print(" ", op.type, op.input, "-->", op.output)
113  print("predict_net:")
114  for op in predict_net.op:
115  print(" ", op.type, op.input, "-->", op.output)
116 
117  with open(args.predict_net, 'wb') as f:
118  f.write(predict_net.SerializeToString())
119  with open(args.init_net, 'wb') as f:
120  f.write(init_net.SerializeToString())
121 
122 if __name__ == "__main__":
123  parser = argparse.ArgumentParser(
124  description="Utilitity to generate Caffe2 benchmark models.")
125  parser.add_argument("operator", help="Caffe2 operator to benchmark.")
126  parser.add_argument("-b", "--blob",
127  help="Instantiate a blob --blob name=dim1,dim2,dim3",
128  action='append')
129  parser.add_argument("--context", help="Context to run on.", default="CPU")
130  parser.add_argument("--kwargs", help="kwargs to pass to operator.",
131  nargs="*", type=parse_kwarg, default=[])
132  parser.add_argument("--init_net", help="Output initialization net.",
133  default="init_net.pb")
134  parser.add_argument("--predict_net", help="Output prediction net.",
135  default="predict_net.pb")
136  parser.add_argument("--benchmark_name",
137  help="Name of the benchmark network",
138  default="benchmark")
139  parser.add_argument("--input_name", help="Name of the input blob.",
140  default="data")
141  parser.add_argument("--output_name", help="Name of the output blob.",
142  default="output")
143  parser.add_argument("--iters",
144  help="Number of iterations to run the operator.",
145  default="1")
146  parser.add_argument("-d", "--debug", help="Print debug information.",
147  action='store_true')
148  parser.add_argument("-c", "--chain",
149  help="Chain ops together (create data dependencies)",
150  action='store_true')
151  args = parser.parse_args()
152  main(args)