Caffe2 - Python API
A deep learning, cross platform ML framework
functional.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 functional
17 # Module caffe2.python.layers.functional
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.python import core, schema, scope, workspace
24 from caffe2.python.layers.layers import (
25  ModelLayer,
26 )
27 import caffe2.proto.caffe2_pb2 as caffe2_pb2
28 import numpy as np
29 import six
30 import logging
31 
32 logger = logging.getLogger(__name__)
33 logger.setLevel(logging.INFO)
34 
35 
36 class Functional(ModelLayer):
37 
38  def __init__(self, model, input_record, output_names_or_num, function,
39  name='functional', output_dtypes=None, **kwargs):
40 
41  # allow coercion
42  input_record = schema.as_record(input_record)
43 
44  super(Functional, self).__init__(model, name, input_record, **kwargs)
45  self._function = function
46  self._kwargs = kwargs
47  return_struct = (
48  isinstance(output_names_or_num, list) or
49  (isinstance(output_names_or_num, six.integer_types) and
50  output_names_or_num != 1)
51  )
52 
53  with scope.NameScope(self.name, reset=True):
54  if isinstance(output_names_or_num, int):
55  struct_output_schema = schema.NewRecord(
56  model.net, schema.RawTuple(output_names_or_num))
57  elif isinstance(output_names_or_num, schema.Field):
58  self.output_schema = output_names_or_num.clone(keep_blobs=True)
59  return
60  else:
61  if not isinstance(output_names_or_num, list):
62  output_names_or_num = [output_names_or_num]
63  out_tuple = [(out, np.void) for out in output_names_or_num]
64  struct_output_schema = schema.NewRecord(
65  model.net, schema.Struct(*out_tuple))
66 
67  num_outputs = len(struct_output_schema.field_blobs())
68 
69  # functional layer returns Struct if more than one outputs or output is
70  # a list, otherwise Scalar
71  if return_struct:
72  self.output_schema = struct_output_schema
73  else:
74  self.output_schema = struct_output_schema[0]
75 
76  # If output_dtypes is provided, use it for output schema. Otherwise
77  # the shape and type will be inferred.
78  if output_dtypes is not None:
79  if not isinstance(output_dtypes, list):
80  output_dtypes = [output_dtypes] * num_outputs
81  assert len(output_dtypes) == num_outputs
82  for dtype, scalar in zip(output_dtypes,
83  self.output_schema.all_scalars()):
84  scalar.set_type(dtype)
85  return
86 
87  # Fake execution of the function to infer shapes and types automatically
88  had_issues = False
89  try:
90  type_net = core.Net('_temp_type_and_shape_inference_net')
91  schema.InitEmptyRecord(type_net, input_record, enforce_types=True)
92 
93  function(type_net, self.input_record, self.output_schema, **kwargs)
94  (shapes, types) = workspace.InferShapesAndTypes([type_net], {})
95  for i in range(num_outputs):
96  scalar_schema = (self.output_schema[i] if return_struct
97  else self.output_schema)
98  blob = scalar_schema()
99  if blob not in types or blob not in shapes:
100  had_issues = True
101  continue
102  if shapes[blob] == []:
103  # Scalar type
104  shape = tuple()
105  elif shapes[blob][0] == 0:
106  shape = tuple(shapes[blob][1:])
107  else:
108  logger.warning("unexpeced shape: {}".format(shapes[blob]))
109  # If batch dimension is not first - give up on shape
110  # inference for that blob
111  had_issues = True
112  continue
113 
114  # TODO(amalevich): Move it to some shared library
115  dtype = None
116  if types[blob] == caffe2_pb2.TensorProto.DOUBLE:
117  dtype = (np.float64, shape)
118  elif types[blob] == caffe2_pb2.TensorProto.FLOAT:
119  dtype = (np.float32, shape)
120  elif types[blob] == caffe2_pb2.TensorProto.INT32:
121  dtype = (np.int32, shape)
122  elif types[blob] == caffe2_pb2.TensorProto.INT64:
123  dtype = (np.int64, shape)
124  elif types[blob] == caffe2_pb2.TensorProto.FLOAT16:
125  dtype = (np.float16, shape)
126 
127  if dtype is not None:
128  scalar_schema.set_type(dtype)
129  except TypeError as ex:
130  had_issues = True
131  logger.warning(str(ex))
132 
133  if had_issues:
134  logger.warning(
135  "Type inference had problems for layer: {}".format(self.name))
136 
137  def add_ops(self, net):
138  self._function(
139  net, self.input_record, self.output_schema, **(self._kwargs))