| | |
| | |
| |
|
| |
|
| |
|
| |
|
| |
|
| | from caffe2.python import core |
| | from caffe2.python.modeling import initializers |
| | from caffe2.python.modeling.parameter_info import ParameterTags |
| |
|
| |
|
| | def _FC_or_packed_FC( |
| | model, op_call, blob_in, blob_out, dim_in, dim_out, weight_init=None, |
| | bias_init=None, WeightInitializer=None, BiasInitializer=None, |
| | enable_tensor_core=False, float16_compute=False, **kwargs |
| | ): |
| | WeightInitializer = initializers.update_initializer( |
| | WeightInitializer, weight_init, ("XavierFill", {}) |
| | ) |
| | BiasInitializer = initializers.update_initializer( |
| | BiasInitializer, bias_init, ("ConstantFill", {}) |
| | ) |
| | if not model.init_params: |
| | WeightInitializer = initializers.ExternalInitializer() |
| | BiasInitializer = initializers.ExternalInitializer() |
| |
|
| | blob_out = blob_out or model.net.NextName() |
| | bias_tags = [ParameterTags.BIAS] |
| | if 'freeze_bias' in kwargs: |
| | bias_tags.append(ParameterTags.COMPUTED_PARAM) |
| |
|
| | weight = model.create_param( |
| | param_name=blob_out + '_w', |
| | shape=[dim_out, dim_in], |
| | initializer=WeightInitializer, |
| | tags=ParameterTags.WEIGHT |
| | ) |
| | bias = model.create_param( |
| | param_name=blob_out + '_b', |
| | shape=[dim_out, ], |
| | initializer=BiasInitializer, |
| | tags=bias_tags |
| | ) |
| |
|
| | |
| | if enable_tensor_core: |
| | kwargs['engine'] = 'TENSORCORE' |
| |
|
| | |
| | if float16_compute: |
| | kwargs['float16_compute'] = True |
| |
|
| | return op_call([blob_in, weight, bias], blob_out, **kwargs) |
| |
|
| |
|
| | def fc(model, *args, **kwargs): |
| | return _FC_or_packed_FC(model, model.net.FC, *args, **kwargs) |
| |
|
| |
|
| | def packed_fc(model, *args, **kwargs): |
| | return _FC_or_packed_FC(model, model.net.PackedFC, *args, **kwargs) |
| |
|
| |
|
| | def fc_decomp( |
| | model, blob_in, blob_out, dim_in, dim_out, |
| | rank_approx=5, weight_init=None, bias_init=None, |
| | WeightInitializer=None, BiasInitializer=None, **kwargs |
| | ): |
| | """FC_Decomp version |
| | Here we assume that the rank of original input is bigger than 5. |
| | """ |
| | WeightInitializer = initializers.update_initializer( |
| | WeightInitializer, weight_init, ("XavierFill", {}) |
| | ) |
| | BiasInitializer = initializers.update_initializer( |
| | BiasInitializer, bias_init, ("ConstantFill", {}) |
| | ) |
| | blob_out = blob_out or model.net.NextName() |
| | u = model.create_param( |
| | param_name=blob_out + '_u', |
| | shape=[dim_out, rank_approx], |
| | initializer=WeightInitializer, |
| | ) |
| | v = model.create_param( |
| | param_name=blob_out + '_v', |
| | shape=[dim_in, rank_approx], |
| | initializer=WeightInitializer, |
| | ) |
| | bias = model.create_param( |
| | param_name=blob_out + '_b', |
| | shape=[dim_out, ], |
| | initializer=BiasInitializer, |
| | ) |
| | return model.net.FC_Decomp([blob_in, u, v, bias], blob_out, **kwargs) |
| |
|
| |
|
| | def fc_prune( |
| | model, blob_in, blob_out, dim_in, dim_out, |
| | weight_init=None, bias_init=None, mask_init=None, |
| | threshold=0.00001, need_compress_rate=False, |
| | comp_lb=0.05, |
| | **kwargs |
| | ): |
| | """FC_Prune version |
| | Runnable so far. Great!:) |
| | """ |
| | weight_init = weight_init if weight_init else ('XavierFill', {}) |
| | bias_init = bias_init if bias_init else ('ConstantFill', {}) |
| | mask_init = mask_init if mask_init else ('ConstantFill', {}) |
| | blob_out = blob_out or model.net.NextName() |
| | compress_rate = blob_out + '_compress_rate' |
| | if model.init_params: |
| | compress_lb = model.param_init_net.ConstantFill( |
| | [], |
| | blob_out + '_lb', |
| | shape=[1], |
| | value=comp_lb |
| | ) |
| | weight = model.param_init_net.__getattr__(weight_init[0])( |
| | [], |
| | blob_out + '_w', |
| | shape=[dim_out, dim_in], |
| | **weight_init[1] |
| | ) |
| | mask = model.param_init_net.ConstantFill( |
| | [], |
| | blob_out + '_m', |
| | shape=[dim_out, dim_in], |
| | value=1.0 |
| | ) |
| | ag_dw = model.param_init_net.__getattr__(mask_init[0])( |
| | [], |
| | blob_out + '_ag_dw', |
| | shape=[dim_out, dim_in], |
| | **mask_init[1] |
| | ) |
| | bias = model.param_init_net.__getattr__(bias_init[0])( |
| | [], |
| | blob_out + '_b', |
| | shape=[dim_out, ], |
| | **bias_init[1] |
| | ) |
| | mask_seq = model.param_init_net.__getattr__(mask_init[0])( |
| | [], |
| | blob_out + '_mask_seq', |
| | shape=[dim_out, dim_in], |
| | **mask_init[1] |
| | ) |
| | thres = model.param_init_net.ConstantFill( |
| | [], |
| | blob_out + '_thres', |
| | shape=[1], |
| | value=threshold |
| | ) |
| | else: |
| | compress_lb = core.ScopedBlobReference( |
| | blob_out + '_lb', model.param_init_net) |
| | weight = core.ScopedBlobReference( |
| | blob_out + '_w', model.param_init_net) |
| | bias = core.ScopedBlobReference( |
| | blob_out + '_b', model.param_init_net) |
| | mask = core.ScopedBlobReference( |
| | blob_out + '_m', model.param_init_net) |
| | ag_dw = core.ScopedBlobReference( |
| | blob_out + '_ag_dw', model.param_init_net) |
| | mask_seq = core.ScopedBlobReference( |
| | blob_out + '_mask_seq', model.param_init_net) |
| | thres = core.ScopedBlobReference( |
| | blob_out + '_thres', model.param_init_net) |
| |
|
| | model.AddParameter(weight) |
| | model.AddParameter(bias) |
| | if need_compress_rate: |
| | return model.net.FC_Prune([blob_in, weight, mask, bias, ag_dw, mask_seq, |
| | thres, compress_lb], |
| | [blob_out, compress_rate], **kwargs) |
| | else: |
| | return model.net.FC_Prune([blob_in, weight, mask, |
| | bias, ag_dw, mask_seq, |
| | thres, compress_lb], |
| | blob_out, **kwargs) |
| |
|
| |
|
| | def fc_sparse( |
| | model, blob_in, blob_out, w_csr, iw, jw, bias, |
| | **kwargs |
| | ): |
| | """FC_Sparse: Only takes in allocated weights""" |
| | if not (w_csr and iw and jw and bias): |
| | print("Warning...") |
| | model.AddParameter(w_csr) |
| | model.AddParameter(iw) |
| | model.AddParameter(jw) |
| | model.AddParameter(bias) |
| | return model.net.FC_Sparse([blob_in, w_csr, iw, jw, bias], |
| | blob_out, **kwargs) |
| |
|
