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