Caffe2 - Python API
A deep learning, cross platform ML framework
predictor_exporter.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 predictor_exporter
17 # Module caffe2.python.predictor.predictor_exporter
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 from caffe2.proto import caffe2_pb2
24 from caffe2.proto import metanet_pb2
25 from caffe2.python import workspace, core, scope
26 from caffe2.python.predictor_constants import predictor_constants
27 import caffe2.python.predictor.serde as serde
29 from builtins import bytes
30 import collections
31 
32 
33 def get_predictor_exporter_helper(submodelNetName):
34  """ constracting stub for the PredictorExportMeta
35  Only used to construct names to subfields,
36  such as calling to predict_net_name
37  Args:
38  submodelNetName - name of the model
39  """
40  stub_net = core.Net(submodelNetName)
41  pred_meta = PredictorExportMeta(predict_net=stub_net,
42  parameters=[],
43  inputs=[],
44  outputs=[],
45  shapes=None,
46  name=submodelNetName,
47  extra_init_net=None)
48  return pred_meta
49 
50 
51 class PredictorExportMeta(collections.namedtuple(
52  'PredictorExportMeta',
53  'predict_net, parameters, inputs, outputs, shapes, name, \
54  extra_init_net, net_type, num_workers')):
55  """
56  Metadata to be used for serializaing a net.
57 
58  parameters, inputs, outputs could be either BlobReference or blob's names
59 
60  predict_net can be either core.Net, NetDef, PlanDef or object
61 
62  Override the named tuple to provide optional name parameter.
63  name will be used to identify multiple prediction nets.
64 
65  net_type is the type field in caffe2 NetDef - can be 'simple', 'dag', etc.
66 
67  num_workers specifies for net type 'dag' how many threads should run ops
68  """
69  def __new__(
70  cls,
71  predict_net,
72  parameters,
73  inputs,
74  outputs,
75  shapes=None,
76  name="",
77  extra_init_net=None,
78  net_type=None,
79  num_workers=None,
80  ):
81  inputs = [str(i) for i in inputs]
82  outputs = [str(o) for o in outputs]
83  assert len(set(inputs)) == len(inputs), (
84  "All inputs to the predictor should be unique")
85  parameters = [str(p) for p in parameters]
86  assert set(parameters).isdisjoint(inputs), (
87  "Parameters and inputs are required to be disjoint. "
88  "Intersection: {}".format(set(parameters).intersection(inputs)))
89  assert set(parameters).isdisjoint(outputs), (
90  "Parameters and outputs are required to be disjoint. "
91  "Intersection: {}".format(set(parameters).intersection(outputs)))
92  shapes = shapes or {}
93 
94  if isinstance(predict_net, (core.Net, core.Plan)):
95  predict_net = predict_net.Proto()
96 
97  assert isinstance(predict_net, (caffe2_pb2.NetDef, caffe2_pb2.PlanDef))
98  return super(PredictorExportMeta, cls).__new__(
99  cls, predict_net, parameters, inputs, outputs, shapes, name,
100  extra_init_net, net_type, num_workers)
101 
102  def inputs_name(self):
103  return utils.get_comp_name(predictor_constants.INPUTS_BLOB_TYPE,
104  self.name)
105 
106  def outputs_name(self):
107  return utils.get_comp_name(predictor_constants.OUTPUTS_BLOB_TYPE,
108  self.name)
109 
110  def parameters_name(self):
111  return utils.get_comp_name(predictor_constants.PARAMETERS_BLOB_TYPE,
112  self.name)
113 
114  def global_init_name(self):
115  return utils.get_comp_name(predictor_constants.GLOBAL_INIT_NET_TYPE,
116  self.name)
117 
118  def predict_init_name(self):
119  return utils.get_comp_name(predictor_constants.PREDICT_INIT_NET_TYPE,
120  self.name)
121 
122  def predict_net_name(self):
123  return utils.get_comp_name(predictor_constants.PREDICT_NET_TYPE,
124  self.name)
125 
126  def train_init_plan_name(self):
127  return utils.get_comp_name(predictor_constants.TRAIN_INIT_PLAN_TYPE,
128  self.name)
129 
130  def train_plan_name(self):
131  return utils.get_comp_name(predictor_constants.TRAIN_PLAN_TYPE,
132  self.name)
133 
134 
135 def prepare_prediction_net(filename, db_type, device_option=None):
136  '''
137  Helper function which loads all required blobs from the db
138  and returns prediction net ready to be used
139  '''
140  metanet_def = load_from_db(filename, db_type, device_option)
141 
142  global_init_net = utils.GetNet(
143  metanet_def, predictor_constants.GLOBAL_INIT_NET_TYPE)
144  workspace.RunNetOnce(global_init_net)
145 
146  predict_init_net = utils.GetNet(
147  metanet_def, predictor_constants.PREDICT_INIT_NET_TYPE)
148  workspace.RunNetOnce(predict_init_net)
149 
150  predict_net = core.Net(
151  utils.GetNet(metanet_def, predictor_constants.PREDICT_NET_TYPE))
152  workspace.CreateNet(predict_net)
153 
154  return predict_net
155 
156 
157 def _global_init_net(predictor_export_meta):
158  net = core.Net("global-init")
159  net.Load(
160  [predictor_constants.PREDICTOR_DBREADER],
161  predictor_export_meta.parameters)
162  net.Proto().external_input.extend([predictor_constants.PREDICTOR_DBREADER])
163  net.Proto().external_output.extend(predictor_export_meta.parameters)
164 
165  # Add the model_id in the predict_net to the global_init_net
166  utils.AddModelIdArg(predictor_export_meta, net.Proto())
167  return net.Proto()
168 
169 
170 def get_meta_net_def(predictor_export_meta, ws=None):
171  """
172  """
173 
174  ws = ws or workspace.C.Workspace.current
175  meta_net_def = metanet_pb2.MetaNetDef()
176 
177  # Predict net is the core network that we use.
178  utils.AddNet(meta_net_def, predictor_export_meta.predict_init_name(),
179  utils.create_predict_init_net(ws, predictor_export_meta))
180  utils.AddNet(meta_net_def, predictor_export_meta.global_init_name(),
181  _global_init_net(predictor_export_meta))
182  utils.AddNet(meta_net_def, predictor_export_meta.predict_net_name(),
183  utils.create_predict_net(predictor_export_meta))
184  utils.AddBlobs(meta_net_def, predictor_export_meta.parameters_name(),
185  predictor_export_meta.parameters)
186  utils.AddBlobs(meta_net_def, predictor_export_meta.inputs_name(),
187  predictor_export_meta.inputs)
188  utils.AddBlobs(meta_net_def, predictor_export_meta.outputs_name(),
189  predictor_export_meta.outputs)
190  return meta_net_def
191 
192 
193 def set_model_info(meta_net_def, project_str, model_class_str, version):
194  assert isinstance(meta_net_def, metanet_pb2.MetaNetDef)
195  meta_net_def.modelInfo.project = project_str
196  meta_net_def.modelInfo.modelClass = model_class_str
197  meta_net_def.modelInfo.version = version
198 
199 
200 def save_to_db(db_type, db_destination, predictor_export_meta):
201  meta_net_def = get_meta_net_def(predictor_export_meta)
202  with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
203  workspace.FeedBlob(
204  predictor_constants.META_NET_DEF,
205  serde.serialize_protobuf_struct(meta_net_def)
206  )
207 
208  blobs_to_save = [predictor_constants.META_NET_DEF] + \
209  predictor_export_meta.parameters
210  op = core.CreateOperator(
211  "Save",
212  blobs_to_save, [],
213  absolute_path=True,
214  db=db_destination, db_type=db_type)
215 
216  workspace.RunOperatorOnce(op)
217 
218 
219 def load_from_db(filename, db_type, device_option=None):
220  # global_init_net in meta_net_def will load parameters from
221  # predictor_constants.PREDICTOR_DBREADER
222  create_db = core.CreateOperator(
223  'CreateDB', [],
224  [core.BlobReference(predictor_constants.PREDICTOR_DBREADER)],
225  db=filename, db_type=db_type)
226  assert workspace.RunOperatorOnce(create_db), (
227  'Failed to create db {}'.format(filename))
228 
229  # predictor_constants.META_NET_DEF is always stored before the parameters
230  load_meta_net_def = core.CreateOperator(
231  'Load',
232  [core.BlobReference(predictor_constants.PREDICTOR_DBREADER)],
233  [core.BlobReference(predictor_constants.META_NET_DEF)])
234  assert workspace.RunOperatorOnce(load_meta_net_def)
235 
236  blob = workspace.FetchBlob(predictor_constants.META_NET_DEF)
237  meta_net_def = serde.deserialize_protobuf_struct(
238  blob if isinstance(blob, bytes)
239  else str(blob).encode('utf-8'),
240  metanet_pb2.MetaNetDef)
241 
242  if device_option is None:
243  device_option = scope.CurrentDeviceScope()
244 
245  if device_option is not None:
246  # Set the device options of all loaded blobs
247  for kv in meta_net_def.nets:
248  net = kv.value
249  for op in net.op:
250  op.device_option.CopyFrom(device_option)
251 
252  return meta_net_def