| |
| |
|
|
|
|
|
|
|
|
|
|
| from caffe2.python.helpers.arg_scope import get_current_scope |
| from caffe2.python import schema |
| from caffe2.python.layers.layers import ModelLayer |
| from caffe2.python.layers.sampling_trainable_mixin import SamplingTrainableMixin |
| import math |
| import numpy as np |
|
|
|
|
| def get_fc_predictor_version(fc_version): |
| assert fc_version in ["fp32", "fp16"], ( |
| "Only support fp32 and fp16 for the fully connected layer " |
| "in the predictor net, the provided FC precision is {}".format(fc_version) |
| ) |
| return fc_version |
|
|
|
|
| class FC(SamplingTrainableMixin, ModelLayer): |
|
|
| def __init__(self, model, input_record, output_dims, weight_init=None, |
| bias_init=None, weight_optim=None, bias_optim=None, name='fc', |
| weight_reg=None, bias_reg=None, clip_param=None, |
| max_fc_size=None, axis=1, transposed=False, |
| uniform_weight_init_scale_numerator=1.0, |
| **kwargs): |
| super(FC, self).__init__(model, name, input_record, **kwargs) |
| assert isinstance(input_record, schema.Scalar), ( |
| "Incorrect input type {}".format(input_record)) |
| assert len(input_record.field_types()[0].shape) > 0, ( |
| "FC expects limited dimensions of the input tensor") |
| assert axis >= 1, "axis {} should >= 1.".format(axis) |
| self.axis = axis |
| input_dims = np.prod(input_record.field_types()[0].shape[axis - 1:]) |
|
|
| assert input_dims > 0, ( |
| "FC expects input dimensions > 0, got {}".format(input_dims)) |
|
|
| self.clip_args = None |
| if (clip_param is not None): |
| assert len(clip_param) == 2, ( |
| 'clip_param must be a tuple / list ' |
| 'of length 2 and in the form of (clip_min, clip max)' |
| ) |
| clip_min, clip_max = clip_param |
| assert clip_min is not None or clip_max is not None, ( |
| 'clip_min, and clip_max in clip_param cannot both be None' |
| ) |
| assert ( |
| (clip_min is None or clip_max is None) or clip_min < clip_max |
| ), ( |
| 'clip_param = [clip_min, clip_max] must have clip_min < clip_max' |
| ) |
| self.clip_args = {} |
| if clip_min is not None: |
| self.clip_args['min'] = clip_min |
| if clip_max is not None: |
| self.clip_args['max'] = clip_max |
|
|
| if uniform_weight_init_scale_numerator is None: |
| uniform_weight_init_scale_numerator = 1.0 |
|
|
| scale = math.sqrt(uniform_weight_init_scale_numerator / input_dims) |
| weight_init = weight_init if weight_init else ( |
| 'UniformFill', {'min': -scale, 'max': scale}) |
| bias_init = bias_init if bias_init else ( |
| 'UniformFill', {'min': -scale, 'max': scale}) |
|
|
| self.output_dim_vec = FC.calculate_fc_output_dims( |
| max_fc_size, input_dims, output_dims) |
|
|
| self.transposed = transposed |
| if self.output_dim_vec is None or len(self.output_dim_vec) == 1: |
| weight_shape = [input_dims, output_dims] if transposed else [output_dims, input_dims] |
| self.w = self.create_param(param_name='w', |
| shape=weight_shape, |
| initializer=weight_init, |
| optimizer=weight_optim, |
| regularizer=weight_reg) |
|
|
| self.b = self.create_param(param_name='b', |
| shape=[output_dims, ], |
| initializer=bias_init, |
| optimizer=bias_optim, |
| regularizer=bias_reg) |
| else: |
| self.w_vec = [] |
| self.b_vec = [] |
|
|
| for idx, output_dim in enumerate(self.output_dim_vec): |
| weight_shape = [input_dims, output_dim] if transposed else [output_dim, input_dims] |
| self.w_vec.append(self.create_param(param_name='w_sub_{}'.format(idx), |
| shape=weight_shape, |
| initializer=weight_init, |
| optimizer=weight_optim, |
| regularizer=weight_reg)) |
|
|
| self.b_vec.append(self.create_param(param_name='b_sub_{}'.format(idx), |
| shape=[output_dim, ], |
| initializer=weight_init, |
| optimizer=weight_optim, |
| regularizer=weight_reg)) |
| if axis == 1: |
| output_shape = (output_dims, ) |
| else: |
| output_shape = list(input_record.field_types()[0].shape)[0: axis - 1] |
| output_shape = tuple(output_shape + [output_dims]) |
|
|
| self.output_schema = schema.Scalar( |
| (np.float32, output_shape), |
| self.get_next_blob_reference('output') |
| ) |
|
|
| @staticmethod |
| def calculate_fc_output_dims(max_fc_size, input_dim, output_dim): |
|
|
| if not max_fc_size or max_fc_size < 0: |
| return None |
|
|
| assert max_fc_size >= input_dim, "Currently we split along the output " \ |
| "dimension. So we need max_fc_size >= input_dim. But, max_fc_size: " \ |
| "{}, input_dim: {}".format(max_fc_size, input_dim) |
|
|
| output_dim_allowed = int(np.floor(max_fc_size / input_dim)) |
| num_fc = int(np.floor((output_dim - 1) / output_dim_allowed) + 1) |
|
|
| output_dim_vec = [output_dim_allowed] * (num_fc - 1) |
|
|
| output_dim_vec.append(output_dim - sum(output_dim_vec)) |
|
|
| return output_dim_vec |
|
|
| def _insert_fc_ops(self, net, params, outputs, version): |
| """ |
| Args: |
| net: the caffe2 net to insert operator |
| params: weight and bias for FC |
| outputs: the output blobs |
| version: support fp32 and fp16 for now. |
| """ |
| if version == "fp32": |
| if self.transposed: |
| return net.FCTransposed( |
| self.input_record.field_blobs() + params, |
| outputs, |
| axis=self.axis, |
| **self.kwargs |
| ) |
| else: |
| return net.FC( |
| self.input_record.field_blobs() + params, |
| outputs, |
| axis=self.axis, |
| **self.kwargs |
| ) |
| elif version == "fp16": |
| return net.FbFCPacked( |
| self.input_record.field_blobs() + params, |
| outputs, |
| axis=self.axis, |
| **self.kwargs |
| ) |
| else: |
| raise Exception("unsupported FC type version {}".format(version)) |
|
|
| def _add_ops(self, net, params, version): |
| """ |
| Args: |
| params : the weight and bias, |
| passed by either add_ops or add_train_ops function |
| version : fp16 or fp32, might support in8 in the future. |
| """ |
| if self.clip_args is not None: |
| clipped_params = [net.NextScopedBlob( |
| 'clipped_%s' % str(p)) for p in params] |
| for p, cp in zip(params, clipped_params): |
| net.Clip([p], [cp], **self.clip_args) |
| params = clipped_params |
|
|
| if self.output_dim_vec is None or len(self.output_dim_vec) == 1: |
| self._insert_fc_ops(net, params, self.output_schema.field_blobs(), version) |
| else: |
| w_vec = params[:int(len(params) / 2)] |
| b_vec = params[int(len(params) / 2):] |
|
|
| assert len(w_vec) == len(b_vec) |
|
|
| output_blob_vec = [] |
|
|
| for i in range(len(self.output_dim_vec)): |
| output_blob = net.NextScopedBlob( |
| 'output_sub_{}'.format(i)) |
| insert_ret = self._insert_fc_ops( |
| net, [w_vec[i], b_vec[i]], [output_blob], version |
| ) |
| output_blob_vec.append(insert_ret) |
| net.Concat(output_blob_vec, |
| self.output_schema.field_blobs() + |
| [self.output_schema.field_blobs()[0] + "_concat_dims"]) |
|
|
| def add_ops(self, net): |
| """Both the predict net and the eval net will call this function |
| """ |
| version_info = get_current_scope().get( |
| get_fc_predictor_version.__name__, {'fc_version': 'fp32'} |
| ) |
| predictor_fc_fp_version = version_info['fc_version'] |
| self._add_ops(net, self.param_blobs, predictor_fc_fp_version) |
|
|
| def add_train_ops(self, net): |
| |
| self._add_ops(net, self.train_param_blobs, "fp32") |
|
|
| def get_fp16_compatible_parameters(self): |
| if self.output_dim_vec is None or len(self.output_dim_vec) == 1: |
| return [self.w] |
| else: |
| return self.w_vec |
|
|
| @property |
| def param_blobs(self): |
| if self.output_dim_vec is None or len(self.output_dim_vec) == 1: |
| return [self.w, self.b] |
| else: |
| return self.w_vec + self.b_vec |
|
|
