| |
| |
|
|
|
|
|
|
|
|
|
|
| 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()) |
|
|
