Caffe2 - Python API
A deep learning, cross platform ML framework
net_construct_bench.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 net_construct_bench
17 # Module caffe2.experiments.python.net_construct_bench
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 import argparse
24 import logging
25 import time
26 
27 from caffe2.python import workspace, data_parallel_model
28 from caffe2.python import cnn
29 
30 import caffe2.python.models.resnet as resnet
31 
32 '''
33 Simple benchmark that creates a data-parallel resnet-50 model
34 and measurs the time.
35 '''
36 
37 
38 logging.basicConfig()
39 log = logging.getLogger("net_construct_bench")
40 log.setLevel(logging.DEBUG)
41 
42 
43 def AddMomentumParameterUpdate(train_model, LR):
44  '''
45  Add the momentum-SGD update.
46  '''
47  params = train_model.GetParams()
48  assert(len(params) > 0)
49  ONE = train_model.param_init_net.ConstantFill(
50  [], "ONE", shape=[1], value=1.0,
51  )
52  NEGONE = train_model.param_init_net.ConstantFill(
53  [], 'NEGONE', shape=[1], value=-1.0,
54  )
55 
56  for param in params:
57  param_grad = train_model.param_to_grad[param]
58  param_momentum = train_model.param_init_net.ConstantFill(
59  [param], param + '_momentum', value=0.0
60  )
61 
62  # Update param_grad and param_momentum in place
63  train_model.net.MomentumSGD(
64  [param_grad, param_momentum, LR],
65  [param_grad, param_momentum],
66  momentum=0.9,
67  nesterov=1
68  )
69 
70  # Update parameters by applying the moment-adjusted gradient
71  train_model.WeightedSum(
72  [param, ONE, param_grad, NEGONE],
73  param
74  )
75 
76 
77 def Create(args):
78  gpus = list(range(args.num_gpus))
79  log.info("Running on gpus: {}".format(gpus))
80 
81  # Create CNNModeLhelper object
82  train_model = cnn.CNNModelHelper(
83  order="NCHW",
84  name="resnet50",
85  use_cudnn=True,
86  cudnn_exhaustive_search=False
87  )
88 
89  # Model building functions
90  def create_resnet50_model_ops(model, loss_scale):
91  [softmax, loss] = resnet.create_resnet50(
92  model,
93  "data",
94  num_input_channels=3,
95  num_labels=1000,
96  label="label",
97  )
98  model.Accuracy([softmax, "label"], "accuracy")
99  return [loss]
100 
101  # SGD
102  def add_parameter_update_ops(model):
103  model.AddWeightDecay(1e-4)
104  ITER = model.Iter("ITER")
105  stepsz = int(30)
106  LR = model.net.LearningRate(
107  [ITER],
108  "LR",
109  base_lr=0.1,
110  policy="step",
111  stepsize=stepsz,
112  gamma=0.1,
113  )
114  AddMomentumParameterUpdate(model, LR)
115 
116  def add_image_input(model):
117  pass
118 
119  start_time = time.time()
120 
121  # Create parallelized model
122  data_parallel_model.Parallelize_GPU(
123  train_model,
124  input_builder_fun=add_image_input,
125  forward_pass_builder_fun=create_resnet50_model_ops,
126  param_update_builder_fun=add_parameter_update_ops,
127  devices=gpus,
128  )
129 
130  ct = time.time() - start_time
131  train_model.net._CheckLookupTables()
132 
133  log.info("Model create for {} gpus took: {} secs".format(len(gpus), ct))
134 
135 
136 def main():
137  # TODO: use argv
138  parser = argparse.ArgumentParser(
139  description="Caffe2: Benchmark for net construction"
140  )
141  parser.add_argument("--num_gpus", type=int, default=1,
142  help="Number of GPUs.")
143  args = parser.parse_args()
144 
145  Create(args)
146 
147 
148 if __name__ == '__main__':
149  workspace.GlobalInit(['caffe2', '--caffe2_log_level=2'])
150  import cProfile
151 
152  cProfile.run('main()', sort="cumulative")