| ## @package sparse_lookup |
| # Module caffe2.python.layers.sparse_lookup |
| |
| |
| |
| |
| |
| from caffe2.python.optimizer import FP16_ENGINES, Optimizer |
| from caffe2.python.helpers.arg_scope import get_current_scope |
| from caffe2.python import schema |
| from caffe2.python.layers.layers import ( |
| get_categorical_limit, |
| get_key, |
| IdList, |
| IdScoreList, |
| IdListWithEvicted, |
| IdScoreListWithEvicted, |
| LayerPsParam, |
| ModelLayer, |
| almost_equal_schemas, |
| ) |
| import collections |
| import functools |
| import logging |
| import math |
| import numpy as np |
| import operator |
| |
| logger = logging.getLogger(__name__) |
| |
| |
| def get_trainer_version_based_on_optim(optim_def): |
| if isinstance(optim_def, Optimizer) and hasattr(optim_def, "engine"): |
| logger.info( |
| "Attempting to set trainer version for engine {}".format(optim_def.engine) |
| ) |
| if optim_def.engine in FP16_ENGINES: |
| logger.info("Setting FP16 trainer for engine {}".format(optim_def.engine)) |
| return "fp16" |
| else: |
| logger.info("Setting FP32 trainer for engine {}".format(optim_def.engine)) |
| return "fp32" |
| else: |
| return "fp32" |
| |
| |
| def get_sparse_lookup_predictor_version( |
| version, |
| blob_size=None, |
| min_blob_size_4bits=None, |
| embedding_dim=None, |
| sparse_feature_name=None, |
| ): |
| assert version in { |
| 'fp32', 'fp16', 'uint8rowwise', 'fused_uint8rowwise', 'fused_uint4rowwise' |
| }, "Unexpected version of sparse_lookup layer {0}".format(version) |
| if version == 'fused_uint4rowwise': |
| if ( |
| blob_size is not None |
| and min_blob_size_4bits is not None |
| and embedding_dim is not None |
| ): |
| if blob_size < min_blob_size_4bits: |
| logger.info( |
| "{} fall back to uint8 because lookup table size {} < min_blob_size_4bits {}".format( |
| sparse_feature_name, |
| blob_size, |
| min_blob_size_4bits, |
| ) |
| ) |
| version = 'fused_uint8rowwise' |
| |
| if embedding_dim % 2 == 1: |
| logger.info( |
| "{} fall back to uint8 because lookup table dimension {} is not divisible by 2".format( |
| sparse_feature_name, embedding_dim |
| ) |
| ) |
| version = 'fused_uint8rowwise' |
| else: |
| raise ValueError( |
| ( |
| "When 4 bit quantization is enabled for {}, " |
| "(i.e., Sparse lookup predictor version:{}), " |
| "requires arguments blob_size:{}, " |
| "min_blob_size_4bits:{}, embedding_dim:{}" |
| ).format( |
| sparse_feature_name, |
| version, |
| blob_size, |
| min_blob_size_4bits, |
| embedding_dim |
| ) |
| ) |
| return version |
| |
| |
| def get_sparse_lookup_trainer_version(version): |
| assert version in {'fp32', 'fp16'},\ |
| "Unexpected version of sparse_lookup layer {0}".format(version) |
| return version |
| |
| def _is_id_list(input_record): |
| return almost_equal_schemas(input_record, IdList) |
| |
| |
| def _is_id_score_list(input_record): |
| return almost_equal_schemas(input_record, |
| IdScoreList, |
| check_field_types=False) |
| |
| |
| class SparseLookup(ModelLayer): |
| _id_list_supported_reducers = [ |
| 'LogMeanExp', 'LogSumExp', 'Max', 'Mean', 'Sum', |
| 'WeightedSum', 'WeightedMean', 'Sqrt', 'None'] |
| |
| _id_score_list_supported_reducers = [ |
| 'PositionWeighted', 'RecencyWeighted', 'Mean', 'Sum', 'WeightedSum', |
| 'WeightedMean', 'None' |
| ] |
| |
| _fp16_compatible_init_op_types = [ |
| 'Float16UniformFill' |
| ] |
| |
| _fp16_compatible_reducers = [ |
| 'Sum', 'Mean', 'Sqrt', 'PositionWeighted', 'RecencyWeighted', |
| ] |
| |
| def __init__(self, model, input_record, inner_shape, reducer, |
| weight_init=None, weight_optim=None, |
| name='sparse_lookup', regularizer=None, use_external_weights=False, |
| uniform_weight_init_scale_numerator=1.0, **kwargs): |
| |
| super(SparseLookup, self).__init__(model, name, input_record, **kwargs) |
| |
| self.sparse_key = get_key(self.input_record)() |
| logger.info("Setup the sparse lookup layer for " + self.sparse_key) |
| |
| # TODO Add some asserts about input type |
| if isinstance(inner_shape, int): |
| inner_shape = [inner_shape] |
| assert isinstance(inner_shape, list) or isinstance(inner_shape, tuple),\ |
| "Unexpected type for inner_shape, expected list or tuple, got {0} for {1}".\ |
| format(type(inner_shape), self.sparse_key) |
| |
| if reducer == "PositionWeighted": |
| assert _is_id_score_list(self.input_record), ( |
| "PositionWeighted only support IdScoreList, but got {} for {}" |
| + "please use PositionWeighted layer to convert IdList " |
| + "to IdScoreList" |
| ).format(repr(self.input_record), self.sparse_key) |
| self.external_weights = self.input_record.values() |
| |
| elif reducer == "RecencyWeighted": |
| assert _is_id_score_list(self.input_record), ( |
| "RecencyWeighted only supports IdScoreList, " |
| "while the sparse feature {} is not.".format(self.sparse_key) |
| ) |
| self.external_weights = self.input_record.values() |
| # TODO: create a new type of reducer with external weights to wrap |
| # this and the above two cases since essentially their input formats |
| # are the same. |
| elif use_external_weights: |
| assert _is_id_score_list(self.input_record), ( |
| "Use_external_weights only supports IdScoreList, " |
| "while the sparse feature {} is not.".format(self.sparse_key) |
| ) |
| assert reducer in ["Sum", "WeightedSum"], ( |
| "Use_external_weights only supports Sum reducer, " |
| "while the reducer is {}.".format(reducer) |
| ) |
| self.external_weights = self.input_record.values() |
| self.reducer = reducer |
| self.use_external_weights = use_external_weights |
| |
| input_dim = get_categorical_limit(self.input_record) |
| assert input_dim > 0, "{} should have categorical limit > 0, but got {}".format( |
| self.sparse_key, input_dim |
| ) |
| |
| self.input_dim = input_dim |
| self.shape = [input_dim] + inner_shape |
| |
| self.trainer_version = get_trainer_version_based_on_optim( |
| weight_optim |
| ) |
| |
| self.uniform_weight_init_scale_numerator = uniform_weight_init_scale_numerator |
| default_init_op = self._get_default_init_op() |
| |
| self.weight_init = weight_init or default_init_op |
| |
| self.evicted_values = None |
| if schema.equal_schemas( |
| self.input_record, IdListWithEvicted |
| ) or schema.equal_schemas( |
| self.input_record, IdScoreListWithEvicted, check_field_types=False |
| ): |
| self.evicted_values = self.input_record._evicted_values |
| |
| # If fp16 is used, make sure fp16 init op is used |
| if self.trainer_version == "fp16": |
| assert self.reducer in self._fp16_compatible_reducers or use_external_weights, ( |
| "Fp16 training is enabled. The reducer specified is not supported. " |
| "Got {}. Supported reducers: {}. Right now, in general, sum, mean, " |
| "positional pooling are supported. Attention is not. Please check " |
| "if there is fp16 trained sparse features using advanced pooling.".format( |
| self.reducer, self._fp16_compatible_reducers) |
| ) |
| |
| # if init op is UniformFill, we replace it directly |
| if self.weight_init[0] == "UniformFill": |
| self.weight_init = ("Float16UniformFill", self.weight_init[1]) |
| assert self.weight_init[0] in self._fp16_compatible_init_op_types, ( |
| "Fp16 training is enabled. Init op for weight parameter must be fp16 " |
| "compatibale. Got {}. Supported ops: {}".format( |
| self.weight_init[0], |
| self._fp16_compatible_init_op_types) |
| ) |
| |
| assert regularizer is None, "Regularizer is not compatible with fp16" |
| |
| if self.input_record.lengths.metadata: |
| avg_length = self.input_record.lengths.metadata.expected_value |
| else: |
| avg_length = None |
| |
| self.w = self.create_param( |
| param_name='w', |
| shape=self.shape, |
| initializer=self.weight_init, |
| optimizer=weight_optim, |
| ps_param=LayerPsParam( |
| sparse_key=self.sparse_key, |
| average_length=avg_length), |
| regularizer=regularizer |
| ) |
| if self.evicted_values: |
| self.reinit_vec = self.create_param( |
| param_name="reinit_vec", |
| shape=inner_shape, |
| initializer=self.weight_init, |
| optimizer=model.NoOptim, |
| regularizer=None, |
| ) |
| |
| self.scale_bias_init = ('ConstantFill', {'value': 0.0}) |
| |
| self.scale_bias = self.create_param( |
| param_name='scale_bias', |
| shape=[], |
| initializer=self.scale_bias_init, |
| optimizer=model.NoOptim, |
| ) |
| |
| self.output_schema = schema.Scalar( |
| (np.float32, inner_shape), |
| self.get_next_blob_reference('output'), |
| ) |
| |
| def get_memory_usage(self): |
| return functools.reduce(operator.mul, self.shape) * 4 |
| |
| def get_fp16_compatible_parameters(self): |
| return [self.w] |
| |
| def support_8bit(self): |
| # Rowwise quantization makes sense only if shape it's 2D matrix with |
| # second dimension >= 8 |
| if len(self.shape) != 2 or self.shape[1] < 8: |
| return False |
| return True |
| |
| def get_8bits_compatible_parameters(self, fused=True): |
| if not self.support_8bit(): |
| return [] |
| if fused: |
| RowwiseQuantized8BitsWeight = collections.namedtuple( |
| 'RowwiseQuantized8BitsWeight', 'w' |
| ) |
| return [RowwiseQuantized8BitsWeight(self.w)] |
| else: |
| RowwiseQuantized8BitsWeight = collections.namedtuple( |
| 'RowwiseQuantized8BitsWeight', 'w, scale_bias' |
| ) |
| return [RowwiseQuantized8BitsWeight(self.w, self.scale_bias)] |
| |
| def _get_default_init_op(self): |
| scale = math.sqrt(self.uniform_weight_init_scale_numerator / self.input_dim) |
| |
| if self.trainer_version == 'fp32': |
| default_weight_init = ('UniformFill', {'min': -scale, 'max': scale}) |
| elif self.trainer_version == 'fp16': |
| default_weight_init = ("Float16UniformFill", {'min': -scale, 'max': scale}) |
| else: |
| raise NotImplementedError( |
| "Train version {} is not currently supported for sparse feature {}".format( |
| trainer_version, self.sparse_key |
| ) |
| ) |
| |
| return default_weight_init |
| |
| def _gather_wrapper(self, net, version, in_indices, out): |
| # Gather can work on all kinds of input data types, and output |
| # data with the same type. Convert the output of Gather to float, |
| # because the follow-up Ops expect fp32. |
| if version == 'fp32': |
| return net.Gather([self.w, in_indices], out) |
| elif version == 'fp16': |
| gathered_w = net.Gather([self.w, in_indices], 'gathered_w') |
| return net.HalfToFloat(gathered_w, out) |
| elif version == 'uint8rowwise': |
| gathered_w = net.Gather([self.w, in_indices], 'gathered_w') |
| gathered_scale_bias = net.Gather( |
| [self.scale_bias, in_indices], |
| 'gathered_scale_bias' |
| ) |
| |
| return net.Rowwise8BitQuantizedToFloat( |
| [gathered_w, gathered_scale_bias], out) |
| elif version == 'fused_uint8rowwise': |
| gathered_w = net.Gather([self.w, in_indices], 'gathered_w') |
| return net.Fused8BitRowwiseQuantizedToFloat(gathered_w, out) |
| elif version == 'fused_uint4rowwise': |
| gathered_w = net.Gather([self.w, in_indices], 'gathered_w') |
| return net.Fused4BitRowwiseQuantizedToFloat(gathered_w, out) |
| |
| else: |
| raise "Unsupported version of operators in SparseLookup " +\ |
| "layer: {0} for sparse feature {1}".format( |
| version, self.sparse_key |
| ) |
| |
| def _sparse_lengths_weighted_reducer( |
| self, |
| in_indices, |
| weights, |
| reducer, |
| net, |
| version, |
| grad_on_weights=0, |
| ): |
| op_input = [ |
| self.w, |
| weights, |
| in_indices, |
| self.input_record.lengths(), |
| ] |
| layer_name = 'SparseLengths' + reducer |
| |
| if version in ['fp32', 'fp16']: |
| # SparseLengths* Ops will accept either fp16 or fp32 embedding |
| # matrix and output fp32 pooled embedding |
| # A special case here is that we need FP16 engine for |
| # SparseLengthsWeightedSum when FP16 embeedings are used for |
| # correct backward updates |
| if reducer == "WeightedSum" and version == "fp16": |
| net.SparseLengthsWeightedSum( |
| op_input, |
| self.output_schema.field_blobs(), |
| grad_on_weights=grad_on_weights, |
| engine='FP16', |
| ) |
| else: |
| net.__getattr__(layer_name)( |
| op_input, |
| self.output_schema.field_blobs(), |
| grad_on_weights=grad_on_weights, |
| ) |
| elif version == 'uint8rowwise': |
| op_input.insert(len(op_input), self.scale_bias) |
| net.__getattr__(layer_name + '8BitsRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint8rowwise': |
| net.__getattr__(layer_name + 'Fused8BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint4rowwise': |
| net.__getattr__(layer_name + 'Fused4BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| else: |
| raise "Unsupported version of operator in SparseLookUp " +\ |
| "layer: {0} for sparse feature {1}".format( |
| version, self.sparse_key |
| ) |
| |
| # deal with sparse features of id_list type |
| def _add_ops_id_list(self, net, version): |
| assert self.reducer in self._id_list_supported_reducers, ( |
| "Unsupported reducer: {} for ID_LIST {}".format( |
| self.reducer, self.sparse_key |
| ) |
| ) |
| if self.reducer in ['Sum', 'Mean', 'WeightedSum', 'WeightedMean']: |
| op_input = [self.w, |
| self.input_record.items(), |
| self.input_record.lengths()] |
| |
| # For id list features, the behaviors of 'Sum' and |
| # 'WeightedSum' are identical, since we can regard the weight on each |
| # id as 1. Similarly, for 'Mean' and 'WeightedMean'. |
| if self.reducer == 'WeightedSum': |
| self.reducer = 'Sum' |
| elif self.reducer == 'WeightedMean': |
| self.reducer = 'Mean' |
| |
| layer_name = 'SparseLengths' + self.reducer |
| if version in ['fp32', 'fp16']: |
| # SparseLengths* Ops will accept either fp16 or fp32 embedding |
| # matrix and output fp32 pooled embedding |
| net.__getattr__(layer_name)( |
| op_input, |
| self.output_schema.field_blobs(), |
| ) |
| elif version == 'uint8rowwise': |
| op_input.insert(len(op_input), self.scale_bias) |
| net.__getattr__(layer_name + '8BitsRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint8rowwise': |
| net.__getattr__(layer_name + 'Fused8BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint4rowwise': |
| net.__getattr__(layer_name + 'Fused4BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| else: |
| raise "Unsupported version of operator in SparseLookUp " +\ |
| "layer: {0} for sparse feature {1}".format( |
| version, self.sparse_key |
| ) |
| |
| elif self.reducer == 'Sqrt': |
| sqrt_weight = net.LengthsToWeights( |
| [self.input_record.lengths()], |
| [net.NextScopedBlob('lengths_sqrt')], |
| power=0.5, |
| ) |
| self._sparse_lengths_weighted_reducer( |
| self.input_record.items(), |
| sqrt_weight, |
| 'WeightedSum', net, version) |
| |
| elif self.reducer == 'None': |
| # Gather operator will gather the embedding for each id of |
| # each IdList. |
| self._gather_wrapper(net, version, self.input_record.items(), |
| self.output_schema.field_blobs()) |
| |
| else: |
| table_rows = self._gather_wrapper( |
| net, version, self.input_record.items(), 'table_rows') |
| |
| segment_ids = net.LengthsToSegmentIds( |
| self.input_record.lengths(), |
| net.NextScopedBlob(self.input_record.lengths() + '_sid')) |
| net.__getattr__('SortedSegmentRange' + self.reducer)( |
| [table_rows, segment_ids], |
| self.output_schema.field_blobs(), |
| ) |
| |
| # deal with sparse features of id_score_list type |
| def _add_ops_id_score_list(self, net, version): |
| assert self.reducer in self._id_score_list_supported_reducers, ( |
| "Unsupported reducer: {} for ID_SCORE_LIST {}".format( |
| self.reducer, self.sparse_key |
| ) |
| ) |
| if self.reducer in ['WeightedSum', 'WeightedMean']: |
| self._sparse_lengths_weighted_reducer( |
| self.input_record.keys(), |
| self.input_record.values(), |
| self.reducer, net, version) |
| |
| elif self.reducer in ['PositionWeighted', 'RecencyWeighted'] or self.use_external_weights: |
| self._sparse_lengths_weighted_reducer( |
| self.input_record.keys(), |
| self.external_weights, |
| 'WeightedSum', net, version, grad_on_weights=1) |
| |
| elif self.reducer in ['Sum', 'Mean']: |
| op_input = [self.w, |
| self.input_record.keys(), |
| self.input_record.lengths()] |
| |
| layer_name = 'SparseLengths' + self.reducer |
| |
| if version in ['fp32', 'fp16']: |
| net.__getattr__(layer_name)( |
| op_input, |
| self.output_schema.field_blobs(), |
| ) |
| elif version == 'uint8rowwise': |
| net.__getattr__(layer_name + '8BitsRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint8rowwise': |
| net.__getattr__(layer_name + 'Fused8BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| elif version == 'fused_uint4rowwise': |
| net.__getattr__(layer_name + 'Fused4BitRowwise')( |
| op_input, self.output_schema.field_blobs()) |
| else: |
| raise "Unsupported version of operator in SparseLookUp " +\ |
| "layer: {0} for sparse feature {1}".format( |
| version, self.sparse_key |
| ) |
| |
| elif self.reducer == 'None': |
| # Gather operator will gather the embedding for each id of |
| # each IdList. |
| self._gather_wrapper(net, version, self.input_record.keys(), |
| self.output_schema.field_blobs()) |
| else: |
| raise "Only Sum, Mean, None are supported for IdScoreList input." +\ |
| "Trying to create with {} for sparse feature {}".format( |
| self.reducer, self.sparse_key |
| ) |
| |
| def _add_ops(self, net, version='fp32', is_train=True): |
| if self.evicted_values and is_train: |
| net.CopyRowsToTensor( |
| [self.w, self.evicted_values.get(), self.reinit_vec], [self.w]) |
| if _is_id_list(self.input_record): |
| self._add_ops_id_list(net, version=version) |
| elif _is_id_score_list(self.input_record): |
| self._add_ops_id_score_list(net, version=version) |
| else: |
| raise "Unsupported input type {0}".format(self.input_record) |
| |
| def add_train_ops(self, net): |
| self._add_ops(net, self.trainer_version, is_train=True) |
| |
| def add_ops(self, net): |
| version_info = get_current_scope().get( |
| get_sparse_lookup_predictor_version.__name__, {'version': 'fp32'} |
| ) |
| lookup_table_blob_size = self.shape[0] * self.shape[1] |
| version = get_sparse_lookup_predictor_version( |
| version_info['version'], |
| blob_size=lookup_table_blob_size, |
| min_blob_size_4bits=( |
| version_info['min_blob_size_4bits'] |
| if 'min_blob_size_4bits' in version_info |
| else None |
| ), |
| embedding_dim=self.shape[1], |
| sparse_feature_name=self.sparse_key, |
| ) |
| |
| # TODO(amalevich): Layer should not be responsible for decision about |
| # quantization. |
| if not self.support_8bit() and version in {'uint8rowwise', |
| 'fused_uint8rowwise', |
| 'fused_uint4rowwise'}: |
| version = 'fp16' |
| |
| self._add_ops(net, version, is_train=False) |
