Caffe2 - Python API
A deep learning, cross platform ML framework
convnet_benchmarks.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 ## @package convnet_benchmarks
17 # Module caffe2.python.convnet_benchmarks
18 from __future__ import absolute_import
19 from __future__ import division
20 from __future__ import print_function
21 from __future__ import unicode_literals
22 """
23 Benchmark for common convnets.
24 
25 Speed on Titan X, with 10 warmup steps and 10 main steps and with different
26 versions of cudnn, are as follows (time reported below is per-batch time,
27 forward / forward+backward):
28 
29  CuDNN V3 CuDNN v4
30 AlexNet 32.5 / 108.0 27.4 / 90.1
31 OverFeat 113.0 / 342.3 91.7 / 276.5
32 Inception 134.5 / 485.8 125.7 / 450.6
33 VGG (batch 64) 200.8 / 650.0 164.1 / 551.7
34 
35 Speed on Inception with varied batch sizes and CuDNN v4 is as follows:
36 
37 Batch Size Speed per batch Speed per image
38  16 22.8 / 72.7 1.43 / 4.54
39  32 38.0 / 127.5 1.19 / 3.98
40  64 67.2 / 233.6 1.05 / 3.65
41 128 125.7 / 450.6 0.98 / 3.52
42 
43 Speed on Tesla M40, which 10 warmup steps and 10 main steps and with cudnn
44 v4, is as follows:
45 
46 AlexNet 68.4 / 218.1
47 OverFeat 210.5 / 630.3
48 Inception 300.2 / 1122.2
49 VGG (batch 64) 405.8 / 1327.7
50 
51 (Note that these numbers involve a "full" backprop, i.e. the gradient
52 with respect to the input image is also computed.)
53 
54 To get the numbers, simply run:
55 
56 for MODEL in AlexNet OverFeat Inception; do
57  PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
58  --batch_size 128 --model $MODEL --forward_only True
59 done
60 for MODEL in AlexNet OverFeat Inception; do
61  PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
62  --batch_size 128 --model $MODEL
63 done
64 PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
65  --batch_size 64 --model VGGA --forward_only True
66 PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
67  --batch_size 64 --model VGGA
68 
69 for BS in 16 32 64 128; do
70  PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
71  --batch_size $BS --model Inception --forward_only True
72  PYTHONPATH=../gen:$PYTHONPATH python convnet_benchmarks.py \
73  --batch_size $BS --model Inception
74 done
75 
76 Note that VGG needs to be run at batch 64 due to memory limit on the backward
77 pass.
78 """
79 
80 import argparse
81 
82 from caffe2.python import brew, cnn, workspace
83 from caffe2.python.model_helper import ModelHelper
84 
85 from caffe2.python.models import resnet
86 import numpy as np
87 
88 def MLP(order, cudnn_ws, mkl):
89  model = ModelHelper(name="benchmark")
90  d = 256
91  depth = 20
92  width = 3
93  for i in range(depth):
94  for j in range(width):
95  current = "fc_{}_{}".format(i, j) if i > 0 else "data"
96  next_ = "fc_{}_{}".format(i + 1, j)
97  brew.fc(
98  model,
99  current, next_,
100  dim_in=d, dim_out=d,
101  weight_init=('XavierFill', {}),
102  bias_init=('XavierFill', {}))
103 
104  brew.sum(model, ["fc_{}_{}".format(depth, j) for j in range(width)], ["sum"])
105  brew.fc(model, "sum", "last",
106  dim_in=d, dim_out=1000,
107  weight_init=('XavierFill', {}),
108  bias_init=('XavierFill', {}))
109  xent = model.LabelCrossEntropy(["last", "label"], "xent")
110  if not mkl:
111  model.AveragedLoss(xent, "loss")
112  return model, d
113 
114 
115 def ResNet50(order, cudnn_ws, mkl):
116  my_arg_scope = {'order': order, 'use_cudnn': True,
117  'cudnn_exhaustive_search': True,
118  'ws_nbytes_limit': str(cudnn_ws)}
119  model = ModelHelper(name="alexnet", arg_scope=my_arg_scope)
120  resnet.create_resnet50(model, "data", 3, 1000, is_test=True,
121  final_avg_kernel=14)
122  return model, 448
123 
124 def AlexNet(order, cudnn_ws, mkl):
125  my_arg_scope = {'order': order, 'use_cudnn': True,
126  'cudnn_exhaustive_search': True,
127  'ws_nbytes_limit': str(cudnn_ws)}
128  model = ModelHelper(name="alexnet", arg_scope=my_arg_scope)
129  conv1 = brew.conv(
130  model,
131  "data",
132  "conv1",
133  3,
134  64,
135  11,
136  ('XavierFill', {}),
137  ('ConstantFill', {}),
138  stride=4,
139  pad=2
140  )
141  relu1 = brew.relu(model, conv1, "conv1")
142  pool1 = brew.max_pool(model, relu1, "pool1", kernel=3, stride=2)
143  conv2 = brew.conv(
144  model,
145  pool1,
146  "conv2",
147  64,
148  192,
149  5,
150  ('XavierFill', {}),
151  ('ConstantFill', {}),
152  pad=2
153  )
154  relu2 = brew.relu(model, conv2, "conv2")
155  pool2 = brew.max_pool(model, relu2, "pool2", kernel=3, stride=2)
156  conv3 = brew.conv(
157  model,
158  pool2,
159  "conv3",
160  192,
161  384,
162  3,
163  ('XavierFill', {}),
164  ('ConstantFill', {}),
165  pad=1
166  )
167  relu3 = brew.relu(model, conv3, "conv3")
168  conv4 = brew.conv(
169  model,
170  relu3,
171  "conv4",
172  384,
173  256,
174  3,
175  ('XavierFill', {}),
176  ('ConstantFill', {}),
177  pad=1
178  )
179  relu4 = brew.relu(model, conv4, "conv4")
180  conv5 = brew.conv(
181  model,
182  relu4,
183  "conv5",
184  256,
185  256,
186  3,
187  ('XavierFill', {}),
188  ('ConstantFill', {}),
189  pad=1
190  )
191  relu5 = brew.relu(model, conv5, "conv5")
192  pool5 = brew.max_pool(model, relu5, "pool5", kernel=3, stride=2)
193  fc6 = brew.fc(
194  model, pool5, "fc6", 256 * 6 * 6, 4096, ('XavierFill', {}),
195  ('ConstantFill', {})
196  )
197  relu6 = brew.relu(model, fc6, "fc6")
198  fc7 = brew.fc(
199  model, relu6, "fc7", 4096, 4096, ('XavierFill', {}), ('ConstantFill', {})
200  )
201  relu7 = brew.relu(model, fc7, "fc7")
202  fc8 = brew.fc(
203  model, relu7, "fc8", 4096, 1000, ('XavierFill', {}), ('ConstantFill', {})
204  )
205  pred = brew.softmax(model, fc8, "pred")
206  xent = model.LabelCrossEntropy([pred, "label"], "xent")
207  if not mkl:
208  loss = model.AveragedLoss(xent, "loss")
209  return model, 224
210 
211 
212 def OverFeat(order, cudnn_ws, mkl):
213  my_arg_scope = {'order': order, 'use_cudnn': True,
214  'cudnn_exhaustive_search': True,
215  'ws_nbytes_limit': str(cudnn_ws)}
216  model = ModelHelper(name='overfeat', arg_scope=my_arg_scope)
217  conv1 = brew.conv(
218  model,
219  "data",
220  "conv1",
221  3,
222  96,
223  11,
224  ('XavierFill', {}),
225  ('ConstantFill', {}),
226  stride=4
227  )
228  relu1 = brew.relu(model, conv1, "conv1")
229  pool1 = brew.max_pool(model, relu1, "pool1", kernel=2, stride=2)
230  conv2 = brew.conv(
231  model, pool1, "conv2", 96, 256, 5, ('XavierFill', {}), ('ConstantFill', {})
232  )
233  relu2 = brew.relu(model, conv2, "conv2")
234  pool2 = brew.max_pool(model, relu2, "pool2", kernel=2, stride=2)
235  conv3 = brew.conv(
236  model,
237  pool2,
238  "conv3",
239  256,
240  512,
241  3,
242  ('XavierFill', {}),
243  ('ConstantFill', {}),
244  pad=1
245  )
246  relu3 = brew.relu(model, conv3, "conv3")
247  conv4 = brew.conv(
248  model,
249  relu3,
250  "conv4",
251  512,
252  1024,
253  3,
254  ('XavierFill', {}),
255  ('ConstantFill', {}),
256  pad=1
257  )
258  relu4 = brew.relu(model, conv4, "conv4")
259  conv5 = brew.conv(
260  model,
261  relu4,
262  "conv5",
263  1024,
264  1024,
265  3,
266  ('XavierFill', {}),
267  ('ConstantFill', {}),
268  pad=1
269  )
270  relu5 = brew.relu(model, conv5, "conv5")
271  pool5 = brew.max_pool(model, relu5, "pool5", kernel=2, stride=2)
272  fc6 = brew.fc(
273  model, pool5, "fc6", 1024 * 6 * 6, 3072, ('XavierFill', {}),
274  ('ConstantFill', {})
275  )
276  relu6 = brew.relu(model, fc6, "fc6")
277  fc7 = brew.fc(
278  model, relu6, "fc7", 3072, 4096, ('XavierFill', {}), ('ConstantFill', {})
279  )
280  relu7 = brew.relu(model, fc7, "fc7")
281  fc8 = brew.fc(
282  model, relu7, "fc8", 4096, 1000, ('XavierFill', {}), ('ConstantFill', {})
283  )
284  pred = brew.softmax(model, fc8, "pred")
285  xent = model.LabelCrossEntropy([pred, "label"], "xent")
286  if not mkl:
287  loss = model.AveragedLoss(xent, "loss")
288  return model, 231
289 
290 
291 def VGGA(order, cudnn_ws, mkl):
292  my_arg_scope = {'order': order, 'use_cudnn': True,
293  'cudnn_exhaustive_search': True,
294  'ws_nbytes_limit': str(cudnn_ws)}
295  model = ModelHelper(name='vgg-a', arg_scope=my_arg_scope)
296  conv1 = brew.conv(
297  model,
298  "data",
299  "conv1",
300  3,
301  64,
302  3,
303  ('XavierFill', {}),
304  ('ConstantFill', {}),
305  pad=1
306  )
307  relu1 = brew.relu(model, conv1, "conv1")
308  pool1 = brew.max_pool(model, relu1, "pool1", kernel=2, stride=2)
309  conv2 = brew.conv(
310  model,
311  pool1,
312  "conv2",
313  64,
314  128,
315  3,
316  ('XavierFill', {}),
317  ('ConstantFill', {}),
318  pad=1
319  )
320  relu2 = brew.relu(model, conv2, "conv2")
321  pool2 = brew.max_pool(model, relu2, "pool2", kernel=2, stride=2)
322  conv3 = brew.conv(
323  model,
324  pool2,
325  "conv3",
326  128,
327  256,
328  3,
329  ('XavierFill', {}),
330  ('ConstantFill', {}),
331  pad=1
332  )
333  relu3 = brew.relu(model, conv3, "conv3")
334  conv4 = brew.conv(
335  model,
336  relu3,
337  "conv4",
338  256,
339  256,
340  3,
341  ('XavierFill', {}),
342  ('ConstantFill', {}),
343  pad=1
344  )
345  relu4 = brew.relu(model, conv4, "conv4")
346  pool4 = brew.max_pool(model, relu4, "pool4", kernel=2, stride=2)
347  conv5 = brew.conv(
348  model,
349  pool4,
350  "conv5",
351  256,
352  512,
353  3,
354  ('XavierFill', {}),
355  ('ConstantFill', {}),
356  pad=1
357  )
358  relu5 = brew.relu(model, conv5, "conv5")
359  conv6 = brew.conv(
360  model,
361  relu5,
362  "conv6",
363  512,
364  512,
365  3,
366  ('XavierFill', {}),
367  ('ConstantFill', {}),
368  pad=1
369  )
370  relu6 = brew.relu(model, conv6, "conv6")
371  pool6 = brew.max_pool(model, relu6, "pool6", kernel=2, stride=2)
372  conv7 = brew.conv(
373  model,
374  pool6,
375  "conv7",
376  512,
377  512,
378  3,
379  ('XavierFill', {}),
380  ('ConstantFill', {}),
381  pad=1
382  )
383  relu7 = brew.relu(model, conv7, "conv7")
384  conv8 = brew.conv(
385  model,
386  relu7,
387  "conv8",
388  512,
389  512,
390  3,
391  ('XavierFill', {}),
392  ('ConstantFill', {}),
393  pad=1
394  )
395  relu8 = brew.relu(model, conv8, "conv8")
396  pool8 = brew.max_pool(model, relu8, "pool8", kernel=2, stride=2)
397 
398  fcix = brew.fc(
399  model, pool8, "fcix", 512 * 7 * 7, 4096, ('XavierFill', {}),
400  ('ConstantFill', {})
401  )
402  reluix = brew.relu(model, fcix, "fcix")
403  fcx = brew.fc(
404  model, reluix, "fcx", 4096, 4096, ('XavierFill', {}), ('ConstantFill', {})
405  )
406  relux = brew.relu(model, fcx, "fcx")
407  fcxi = brew.fc(
408  model, relux, "fcxi", 4096, 1000, ('XavierFill', {}), ('ConstantFill', {})
409  )
410  pred = brew.softmax(model, fcxi, "pred")
411  xent = model.LabelCrossEntropy([pred, "label"], "xent")
412  if not mkl:
413  loss = model.AveragedLoss(xent, "loss")
414  return model, 231
415 
416 
417 def _InceptionModule(
418  model, input_blob, input_depth, output_name, conv1_depth, conv3_depths,
419  conv5_depths, pool_depth
420 ):
421  # path 1: 1x1 conv
422  conv1 = brew.conv(
423  model, input_blob, output_name + ":conv1", input_depth, conv1_depth, 1,
424  ('XavierFill', {}), ('ConstantFill', {})
425  )
426  conv1 = brew.relu(model, conv1, conv1)
427  # path 2: 1x1 conv + 3x3 conv
428  conv3_reduce = brew.conv(
429  model, input_blob, output_name + ":conv3_reduce", input_depth,
430  conv3_depths[0], 1, ('XavierFill', {}), ('ConstantFill', {})
431  )
432  conv3_reduce = brew.relu(model, conv3_reduce, conv3_reduce)
433  conv3 = brew.conv(
434  model,
435  conv3_reduce,
436  output_name + ":conv3",
437  conv3_depths[0],
438  conv3_depths[1],
439  3,
440  ('XavierFill', {}),
441  ('ConstantFill', {}),
442  pad=1
443  )
444  conv3 = brew.relu(model, conv3, conv3)
445  # path 3: 1x1 conv + 5x5 conv
446  conv5_reduce = brew.conv(
447  model, input_blob, output_name + ":conv5_reduce", input_depth,
448  conv5_depths[0], 1, ('XavierFill', {}), ('ConstantFill', {})
449  )
450  conv5_reduce = brew.relu(model, conv5_reduce, conv5_reduce)
451  conv5 = brew.conv(
452  model,
453  conv5_reduce,
454  output_name + ":conv5",
455  conv5_depths[0],
456  conv5_depths[1],
457  5,
458  ('XavierFill', {}),
459  ('ConstantFill', {}),
460  pad=2
461  )
462  conv5 = brew.relu(model, conv5, conv5)
463  # path 4: pool + 1x1 conv
464  pool = brew.max_pool(
465  model,
466  input_blob,
467  output_name + ":pool",
468  kernel=3,
469  stride=1,
470  pad=1
471  )
472  pool_proj = brew.conv(
473  model, pool, output_name + ":pool_proj", input_depth, pool_depth, 1,
474  ('XavierFill', {}), ('ConstantFill', {})
475  )
476  pool_proj = brew.relu(model, pool_proj, pool_proj)
477  output = brew.concat(model, [conv1, conv3, conv5, pool_proj], output_name)
478  return output
479 
480 
481 def Inception(order, cudnn_ws, mkl):
482  my_arg_scope = {'order': order, 'use_cudnn': True,
483  'cudnn_exhaustive_search': True,
484  'ws_nbytes_limit': str(cudnn_ws)}
485  model = ModelHelper(name="inception", arg_scope=my_arg_scope)
486  conv1 = brew.conv(
487  model,
488  "data",
489  "conv1",
490  3,
491  64,
492  7,
493  ('XavierFill', {}),
494  ('ConstantFill', {}),
495  stride=2,
496  pad=3
497  )
498  relu1 = brew.relu(model, conv1, "conv1")
499  pool1 = brew.max_pool(model, relu1, "pool1", kernel=3, stride=2, pad=1)
500  conv2a = brew.conv(
501  model, pool1, "conv2a", 64, 64, 1,
502  ('XavierFill', {}), ('ConstantFill', {})
503  )
504  conv2a = brew.relu(model, conv2a, conv2a)
505  conv2 = brew.conv(
506  model,
507  conv2a,
508  "conv2",
509  64,
510  192,
511  3,
512  ('XavierFill', {}),
513  ('ConstantFill', {}),
514  pad=1
515  )
516  relu2 = brew.relu(model, conv2, "conv2")
517  pool2 = brew.max_pool(model, relu2, "pool2", kernel=3, stride=2, pad=1)
518  # Inception modules
519  inc3 = _InceptionModule(
520  model, pool2, 192, "inc3", 64, [96, 128], [16, 32], 32
521  )
522  inc4 = _InceptionModule(
523  model, inc3, 256, "inc4", 128, [128, 192], [32, 96], 64
524  )
525  pool5 = brew.max_pool(model, inc4, "pool5", kernel=3, stride=2, pad=1)
526  inc5 = _InceptionModule(
527  model, pool5, 480, "inc5", 192, [96, 208], [16, 48], 64
528  )
529  inc6 = _InceptionModule(
530  model, inc5, 512, "inc6", 160, [112, 224], [24, 64], 64
531  )
532  inc7 = _InceptionModule(
533  model, inc6, 512, "inc7", 128, [128, 256], [24, 64], 64
534  )
535  inc8 = _InceptionModule(
536  model, inc7, 512, "inc8", 112, [144, 288], [32, 64], 64
537  )
538  inc9 = _InceptionModule(
539  model, inc8, 528, "inc9", 256, [160, 320], [32, 128], 128
540  )
541  pool9 = brew.max_pool(model, inc9, "pool9", kernel=3, stride=2, pad=1)
542  inc10 = _InceptionModule(
543  model, pool9, 832, "inc10", 256, [160, 320], [32, 128], 128
544  )
545  inc11 = _InceptionModule(
546  model, inc10, 832, "inc11", 384, [192, 384], [48, 128], 128
547  )
548  pool11 = brew.average_pool(model, inc11, "pool11", kernel=7, stride=1)
549  fc = brew.fc(
550  model, pool11, "fc", 1024, 1000,
551  ('XavierFill', {}), ('ConstantFill', {})
552  )
553  # It seems that Soumith's benchmark does not have softmax on top
554  # for Inception. We will add it anyway so we can have a proper
555  # backward pass.
556  pred = brew.softmax(model, fc, "pred")
557  xent = model.LabelCrossEntropy([pred, "label"], "xent")
558  if not mkl:
559  loss = model.AveragedLoss(xent, "loss")
560  return model, 224
561 
562 
563 def AddParameterUpdate(model):
564  """ Simple plain SGD update -- not tuned to actually train the models """
565  ITER = brew.iter(model, "iter")
566  LR = model.LearningRate(
567  ITER, "LR", base_lr=-1e-8, policy="step", stepsize=10000, gamma=0.999)
568  ONE = model.param_init_net.ConstantFill([], "ONE", shape=[1], value=1.0)
569  for param in model.params:
570  param_grad = model.param_to_grad[param]
571  model.WeightedSum([param, ONE, param_grad, LR], param)
572 
573 
574 def Benchmark(model_gen, arg):
575  model, input_size = model_gen(arg.order, arg.cudnn_ws, arg.mkl)
576  model.Proto().type = arg.net_type
577  model.Proto().num_workers = arg.num_workers
578 
579  # In order to be able to run everything without feeding more stuff, let's
580  # add the data and label blobs to the parameter initialization net as well.
581  if arg.order == "NCHW":
582  input_shape = [arg.batch_size, 3, input_size, input_size]
583  else:
584  input_shape = [arg.batch_size, input_size, input_size, 3]
585  if arg.model == "MLP":
586  input_shape = [arg.batch_size, input_size]
587 
588  model.param_init_net.GaussianFill(
589  [],
590  "data",
591  shape=input_shape,
592  mean=0.0,
593  std=1.0
594  )
595  #MKL doesn't support int, so have to use numpy
596  if arg.mkl:
597  label = np.random.randint(low=0, high=1000, size=(arg.batch_size,)).astype(np.int32)
598  workspace.FeedBlob("label", label)
599  else:
600  model.param_init_net.UniformIntFill(
601  [],
602  "label",
603  shape=[arg.batch_size, ],
604  min=0,
605  max=999
606  )
607 
608  if arg.forward_only:
609  print('{}: running forward only.'.format(arg.model))
610  else:
611  if arg.mkl:
612  print(
613  '==WARNING==\n'
614  'forward-backward not supported yet in MKL, so exiting'
615  )
616  print('{}: running forward-backward.'.format(arg.model))
617  model.AddGradientOperators(["loss"])
618  AddParameterUpdate(model)
619  if arg.order == 'NHWC':
620  print(
621  '==WARNING==\n'
622  'NHWC order with CuDNN may not be supported yet, so I might\n'
623  'exit suddenly.'
624  )
625 
626  if not arg.cpu:
627  if arg.mkl:
628  model.param_init_net.RunAllOnMKL()
629  model.net.RunAllOnMKL()
630  else:
631  model.param_init_net.RunAllOnGPU()
632  model.net.RunAllOnGPU()
633 
634  if arg.engine:
635  for op in model.net.Proto().op:
636  op.engine = arg.engine
637 
638  if arg.dump_model:
639  # Writes out the pbtxt for benchmarks on e.g. Android
640  with open(
641  "{0}_init_batch_{1}.pbtxt".format(arg.model, arg.batch_size), "w"
642  ) as fid:
643  fid.write(str(model.param_init_net.Proto()))
644  with open("{0}.pbtxt".format(arg.model, arg.batch_size), "w") as fid:
645  fid.write(str(model.net.Proto()))
646 
647  workspace.RunNetOnce(model.param_init_net)
648  workspace.CreateNet(model.net)
649  workspace.BenchmarkNet(
650  model.net.Proto().name, arg.warmup_iterations, arg.iterations,
651  arg.layer_wise_benchmark)
652 
653 
654 def GetArgumentParser():
655  parser = argparse.ArgumentParser(description="Caffe2 benchmark.")
656  parser.add_argument(
657  "--batch_size",
658  type=int,
659  default=128,
660  help="The batch size."
661  )
662  parser.add_argument("--model", type=str, help="The model to benchmark.")
663  parser.add_argument(
664  "--order",
665  type=str,
666  default="NCHW",
667  help="The order to evaluate."
668  )
669  parser.add_argument(
670  "--cudnn_ws",
671  type=int,
672  help="The cudnn workspace size."
673  )
674  parser.add_argument(
675  "--iterations",
676  type=int,
677  default=10,
678  help="Number of iterations to run the network."
679  )
680  parser.add_argument(
681  "--warmup_iterations",
682  type=int,
683  default=10,
684  help="Number of warm-up iterations before benchmarking."
685  )
686  parser.add_argument(
687  "--forward_only",
688  action='store_true',
689  help="If set, only run the forward pass."
690  )
691  parser.add_argument(
692  "--layer_wise_benchmark",
693  action='store_true',
694  help="If True, run the layer-wise benchmark as well."
695  )
696  parser.add_argument(
697  "--cpu",
698  action='store_true',
699  help="If True, run testing on CPU instead of GPU."
700  )
701  parser.add_argument(
702  "--mkl",
703  action='store_true',
704  help="If True, run testing on CPU-MKL instead of GPU."
705  )
706  parser.add_argument(
707  "--engine",
708  type=str,
709  default="",
710  help="If set, blindly prefer the given engine(s) for every op.")
711  parser.add_argument(
712  "--dump_model",
713  action='store_true',
714  help="If True, dump the model prototxts to disk."
715  )
716  parser.add_argument("--net_type", type=str, default="simple")
717  parser.add_argument("--num_workers", type=int, default=2)
718  parser.add_argument("--use-nvtx", default=False, action='store_true')
719  parser.add_argument("--htrace_span_log_path", type=str)
720  return parser
721 
722 
723 if __name__ == '__main__':
724  args, extra_args = GetArgumentParser().parse_known_args()
725  if (
726  not args.batch_size or not args.model or not args.order
727  ):
728  GetArgumentParser().print_help()
729  else:
730  workspace.GlobalInit(
731  ['caffe2', '--caffe2_log_level=0'] + extra_args +
732  (['--caffe2_use_nvtx'] if args.use_nvtx else []) +
733  (['--caffe2_htrace_span_log_path=' + args.htrace_span_log_path]
734  if args.htrace_span_log_path else []))
735 
736  model_map = {
737  'AlexNet': AlexNet,
738  'OverFeat': OverFeat,
739  'VGGA': VGGA,
740  'Inception': Inception,
741  'ResNet50': ResNet50,
742  'MLP': MLP,
743  }
744  Benchmark(model_map[args.model], args)