Caffe2 - Python API
A deep learning, cross platform ML framework
build_index.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 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 from __future__ import unicode_literals
20 
21 import numpy as np
22 
23 from caffe2.python import core, schema
24 from caffe2.python.layers.layers import ModelLayer
25 
26 
28  """
29  This layer aims to build a mapping from raw keys to indices within [0, max_index).
30  The mapping is continuously built during training. The mapping will be frozen during
31  evaluation and prediction. Unseen keys will be assigned to index 0.
32  """
33 
34  def __init__(
35  self, model,
36  input_record,
37  max_index,
38  name='map_to_range',
39  **kwargs
40  ):
41  super(MapToRange, self).__init__(model, name, input_record, **kwargs)
42 
43  assert max_index > 0
44  assert isinstance(input_record, schema.Scalar)
45 
46  self.max_index = max_index
47 
48  self.handler = self.create_param(
49  param_name='handler',
50  shape=None,
51  initializer=('LongIndexCreate', {'max_elements': self.max_index}),
52  optimizer=model.NoOptim
53  )
54 
56  ('indices', schema.Scalar(
57  np.int64, self.get_next_blob_reference("indices")
58  )),
59  ('handler', schema.Scalar(
60  np.void, self.handler
61  )),
62  )
63 
64  def add_train_ops(self, net):
65  if self.input_record.field_type().base != np.int64:
66  keys = net.Cast(
67  self.input_record(),
68  net.NextScopedBlob("indices_before_mapping"),
69  to=core.DataType.INT64
70  )
71  else:
72  keys = self.input_record()
73 
74  # Load keys into indices
75  indices = net.IndexGet([self.handler, keys],
76  self.output_schema.indices())
77 
78  net.StopGradient(indices, indices)
79 
80  def add_eval_ops(self, net):
81  net.IndexFreeze(self.handler, self.handler)
82  self.add_train_ops(net)
83 
84  def add_ops(self, net):
85  self.add_eval_ops(net)
def get_next_blob_reference(self, name)
Definition: layers.py:352
def create_param(self, param_name, shape, initializer, optimizer, ps_param=None, regularizer=None)
Definition: layers.py:337