1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
|
## @package bucket_weighted
# Module caffe2.python.layers.bucket_weighted
import logging
import numpy as np
from caffe2.python import core, schema
from caffe2.python.layers.layers import (
get_categorical_limit,
ModelLayer,
)
from caffe2.python.layers.tags import Tags
logger = logging.getLogger(__name__)
class BucketWeighted(ModelLayer):
def __init__(self, model, input_record, max_score=0, bucket_boundaries=None,
hash_buckets=True, weight_optim=None, name="bucket_weighted"):
super(BucketWeighted, self).__init__(model, name, input_record)
assert isinstance(input_record, schema.List), "Incorrect input type"
self.bucket_boundaries = bucket_boundaries
self.hash_buckets = hash_buckets
if bucket_boundaries is not None:
self.shape = len(bucket_boundaries) + 1
elif max_score > 0:
self.shape = max_score
else:
self.shape = get_categorical_limit(input_record)
self.bucket_w = self.create_param(param_name='bucket_w',
shape=[self.shape, ],
initializer=('ConstantFill', {'value': 1.0}),
optimizer=weight_optim)
self.output_schema = schema.Struct(
('bucket_weights',
schema.Scalar((np.float32, self.shape),
self.get_next_blob_reference("bucket_w_gather")))
)
self.tags.update({Tags.HANDLE_AS_SPARSE_LAYER})
def get_memory_usage(self):
return self.shape
def add_ops(self, net):
if self.bucket_boundaries is not None:
buckets_int = net.Bucketize(
self.input_record.values(),
"buckets_int",
boundaries=self.bucket_boundaries
)
else:
buckets = self.input_record.values()
buckets_int = net.Cast(
buckets,
"buckets_int",
to=core.DataType.INT32
)
if self.hash_buckets:
buckets_int = net.IndexHash(
buckets_int, "hashed_buckets_int", seed=0, modulo=self.shape
)
net.Gather(
[self.bucket_w, buckets_int],
self.output_schema.bucket_weights.field_blobs())
|
