def __init__(self,
signed: bool,
narrow_range: bool,
quant_type: QuantType,
ms_bit_width_to_clamp: int,
clamp_at_least_init_val: bool,
min_overall_bit_width: Optional[int],
max_overall_bit_width: Optional[int],
msb_clamp_bit_width_impl_type: BitWidthImplType,
override_pretrained_bit_width: bool):
super(ClampQuantProxy, self).__init__()
if quant_type == QuantType.FP:
self.tensor_quant = IdentityPrescaledIntQuant()
elif quant_type == QuantType.INT:
msb_clamp_bit_width_impl = MsbClampParameterBitWidth(ms_bit_width_to_clamp=ms_bit_width_to_clamp,
clamp_at_least_init_val=clamp_at_least_init_val,
min_overall_bit_width=min_overall_bit_width,
max_overall_bit_width=max_overall_bit_width,
bit_width_impl_type=msb_clamp_bit_width_impl_type,
override_pretrained=override_pretrained_bit_width)
tensor_clamp_impl = TensorClamp()
float_to_int_impl = RestrictValue(restrict_value_type=RestrictValueType.INT,
float_to_int_impl_type=FloatToIntImplType.ROUND,
min_val=None)
tensor_quant_impl = PrescaledRestrictIntQuantWithInputBitWidth
self.tensor_quant = tensor_quant_impl(signed=signed,
narrow_range=narrow_range,
tensor_clamp_impl=tensor_clamp_impl,
float_to_int_impl=float_to_int_impl,
msb_clamp_bit_width_impl=msb_clamp_bit_width_impl)
else:
raise Exception("Quantization type {} not supported for accumulators.".format(quant_type))
def __init__(self,
quant_type: QuantType,
bit_width: Optional[int],
narrow_range: bool) -> None:
super(BiasQuantProxy, self).__init__()
if quant_type == QuantType.FP:
self.tensor_quant = None
elif quant_type == QuantType.INT:
tensor_clamp_impl = TensorClamp()
float_to_int_impl = RestrictValue(restrict_value_type=RestrictValueType.INT,
float_to_int_impl_type=FloatToIntImplType.ROUND,
min_val=None)
if bit_width is not None:
bit_width_impl = BitWidthConst(bit_width, restrict_bit_width_type=RestrictValueType.INT)
self.tensor_quant = PrescaledRestrictIntQuant(narrow_range=narrow_range,
signed=True,
tensor_clamp_impl=tensor_clamp_impl,
msb_clamp_bit_width_impl=bit_width_impl,
float_to_int_impl=float_to_int_impl)
self.requires_input_bit_width = False
else:
self.tensor_quant = PrescaledIntQuant(narrow_range=narrow_range,
signed=True,
tensor_clamp_impl=tensor_clamp_impl,
float_to_int_impl=float_to_int_impl)
self.requires_input_bit_width = True
else:
raise Exception('Quantization type {} not supported for bias quant.'
.format(str(quant_type)))
def __init__(self,
signed: bool,
quant_type: QuantType,
ls_bit_width_to_trunc: int,
trunc_at_least_init_val: bool,
min_overall_bit_width: Optional[int],
max_overall_bit_width: Optional[int],
lsb_trunc_bit_width_impl_type: BitWidthImplType,
explicit_rescaling: bool,
override_pretrained_bit_width: bool):
super(TruncQuantProxy, self).__init__()
self.register_buffer(ZERO_HW_SENTINEL_NAME, torch.tensor(ZERO_HW_SENTINEL_VALUE))
self.explicit_rescaling = explicit_rescaling
if quant_type == QuantType.FP:
self.lsb_trunc_bit_width_impl = ZeroLsbTruncBitWidth()
self.tensor_quant = IdentityPrescaledIntQuant()
elif quant_type == QuantType.INT:
self.lsb_trunc_bit_width_impl = LsbTruncParameterBitWidth(ls_bit_width_to_trunc=ls_bit_width_to_trunc,
trunc_at_least_init_val=trunc_at_least_init_val,
min_overall_bit_width=min_overall_bit_width,
max_overall_bit_width=max_overall_bit_width,
bit_width_impl_type=lsb_trunc_bit_width_impl_type,
override_pretrained=override_pretrained_bit_width)
tensor_clamp_impl = TensorClamp()
float_to_int_impl = RestrictValue(restrict_value_type=RestrictValueType.INT,
float_to_int_impl_type=FloatToIntImplType.FLOOR,
min_val=None)
self.tensor_quant = PrescaledIntQuant(signed=signed,
narrow_range=False,
tensor_clamp_impl=tensor_clamp_impl,
float_to_int_impl=float_to_int_impl)
else:
raise Exception("Quantization type {} not supported for accumulators.".format(quant_type))
def __init__(self,
num_features,
eps=1e-5,
momentum=0.1,
restrict_value_type: RestrictValueType = RestrictValueType.FP,
impl_type: ScalingImplType = ScalingImplType.STATS,
bias_quant_type: QuantType = QuantType.FP,
bias_narrow_range: bool = False,
bias_bit_width: int = None):
QuantLayer.__init__(self,
compute_output_scale=False,
compute_output_bit_width=False,
return_quant_tensor=False)
nn.Module.__init__(self)
if bias_quant_type != QuantType.FP and not (self.compute_output_scale and self.compute_output_bit_width):
raise Exception("Quantizing bias requires to compute output scale and output bit width")
self.weight = nn.Parameter(torch.ones(num_features))
self.bias = nn.Parameter(torch.zeros(num_features))
if impl_type == ScalingImplType.PARAMETER_FROM_STATS:
self.running_mean = None
self.running_var = None
elif impl_type == ScalingImplType.STATS:
self.register_buffer('running_mean', torch.zeros(num_features))
self.register_buffer('running_var', torch.ones(num_features))
else:
raise Exception("Scaling mode not supported")
self.eps = eps
self.momentum = momentum
self.impl_type = impl_type
self.num_features = num_features
self.restrict_value_type = restrict_value_type
self.bias_quant = BiasQuantProxy(quant_type=bias_quant_type,
narrow_range=bias_narrow_range,
bit_width=bias_bit_width)
self.restrict_weight = RestrictValue(restrict_value_type=restrict_value_type,
float_to_int_impl_type=FloatToIntImplType.ROUND,
min_val=None)
self.restrict_scaling_preprocess = RestrictValue.restrict_value_op(restrict_value_type,
restrict_value_op_impl_type=
RestrictValueOpImplType.TORCH_MODULE)
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs):
weight_key = prefix + 'weight'
bias_key = prefix + 'bias'
running_mean_key = prefix + 'running_mean'
running_var_key = prefix + 'running_var'
num_batches_tracked_key = prefix + 'num_batches_tracked'
# If it's converting a FP BN into weight/bias impl
if self.impl_type == ScalingImplType.PARAMETER_FROM_STATS \
and running_mean_key in state_dict and running_var_key in state_dict:
weight_init, bias_init = mul_add_from_bn(bn_bias=state_dict[bias_key],
bn_weight=state_dict[weight_key],
bn_mean=state_dict[running_mean_key],
bn_var=state_dict[running_var_key],
bn_eps=self.eps,
affine_only=False)
restrict_op = RestrictValue.restrict_value_op(restrict_value_type=self.restrict_value_type,
restrict_value_op_impl_type=RestrictValueOpImplType.TORCH_FN)
self.weight_sign = torch.sign(weight_init.data)
weight_init = weight_init.detach().clone().abs().data
self.weight.data = restrict_op(weight_init)
self.bias.data = bias_init.detach().clone().data
del state_dict[bias_key]
del state_dict[weight_key]
del state_dict[running_mean_key]
del state_dict[running_var_key]
del state_dict[num_batches_tracked_key]
super(QuantBatchNorm2d, self)._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
if config.IGNORE_MISSING_KEYS and bias_key in missing_keys:
missing_keys.remove(bias_key)
if config.IGNORE_MISSING_KEYS and weight_key in missing_keys:
missing_keys.remove(weight_key)
if num_batches_tracked_key in unexpected_keys:
unexpected_keys.remove(num_batches_tracked_key)
def _weight_quant_init_impl(bit_width: Optional[int],
quant_type: QuantType,
narrow_range: bool,
scaling_override: Optional[nn.Module],
restrict_scaling_type: RestrictValueType,
scaling_const: float,
scaling_stats_op: StatsOp,
scaling_impl_type: ScalingImplType,
scaling_stats_reduce_dim: Optional[int],
scaling_shape: Tuple[int, ...],
scaling_min_val: Optional[float],
bit_width_impl_type: Optional[BitWidthImplType],
restrict_bit_width_type: Optional[RestrictValueType],
min_overall_bit_width: Optional[int],
max_overall_bit_width: Optional[int],
bit_width_impl_override: Optional[Union[BitWidthConst, BitWidthParameter]],
scaling_stats_input_view_shape_impl: StatsInputViewShapeImpl,
scaling_stats_input_concat_dim: int,
ternary_threshold: Optional[float],
scaling_stats_sigma: Optional[float],
tracked_parameter_list: List[torch.nn.Parameter],
zero_hw_sentinel: torch.Tensor,
override_pretrained_bit_width: bool):
if quant_type == QuantType.FP:
tensor_quant = IdentityQuant()
else:
if scaling_impl_type != ScalingImplType.OVERRIDE and scaling_override is not None:
raise Exception("Overriding scaling requires to set ScalingImplType to OVERRIDE explicitly.")
if scaling_impl_type == ScalingImplType.OVERRIDE and scaling_override is None:
raise Exception("Overriding scaling requires to pass a scaling impl module.")
if scaling_impl_type == ScalingImplType.OVERRIDE and scaling_override is not None:
scaling_impl = scaling_override
elif scaling_impl_type == ScalingImplType.STATS \
or scaling_impl_type == ScalingImplType.AFFINE_STATS \
or scaling_impl_type == ScalingImplType.PARAMETER_FROM_STATS:
stats_scaling = ParameterStatsScaling(stats_op=scaling_stats_op,
restrict_scaling_type=restrict_scaling_type,
tracked_parameter_list=tracked_parameter_list,
stats_input_view_shape_impl=scaling_stats_input_view_shape_impl,
stats_input_concat_dim=scaling_stats_input_concat_dim,
sigma=scaling_stats_sigma,
scaling_min_val=scaling_min_val,
stats_reduce_dim=scaling_stats_reduce_dim,
stats_output_shape=scaling_shape,
affine=scaling_impl_type == ScalingImplType.AFFINE_STATS)
if scaling_impl_type == ScalingImplType.PARAMETER_FROM_STATS:
if quant_type == QuantType.BINARY or quant_type == QuantType.TERNARY:
raise Exception("Parameter from stats scaling is currently not supported for binary/ternary")
scaling_init = stats_scaling(zero_hw_sentinel).detach()
scaling_impl = StandaloneScaling(scaling_init=scaling_init,
parameter_shape=scaling_shape,
restrict_scaling_type=restrict_scaling_type,
is_parameter=True,
scaling_min_val=scaling_min_val)
else:
scaling_impl = stats_scaling
elif scaling_impl_type == ScalingImplType.CONST or scaling_impl_type == ScalingImplType.HE:
if scaling_impl_type == ScalingImplType.HE:
scaling_const = 0.0
for param in tracked_parameter_list: # takes average of He scaling over parameter list
two_dim_param = param.view(param.shape[0], -1)
scaling_const += math.sqrt(2.0 / two_dim_param.shape[1])
scaling_const /= len(tracked_parameter_list)
scaling_init = torch.tensor(scaling_const)
scaling_impl = StandaloneScaling(scaling_init=scaling_init,
parameter_shape=SCALING_SCALAR_SHAPE,
restrict_scaling_type=restrict_scaling_type,
is_parameter=False,
scaling_min_val=None)
else:
raise Exception("Scaling type {} not supported for weight quantization"
.format(str(scaling_impl_type)))
if bit_width == 1 and quant_type == QuantType.BINARY:
tensor_quant = BinaryQuant(scaling_impl=scaling_impl)
elif bit_width == 2 and quant_type == QuantType.TERNARY:
tensor_quant = TernaryQuant(scaling_impl=scaling_impl, threshold=ternary_threshold)
elif bit_width >= 2 and quant_type == QuantType.INT:
if bit_width_impl_override is None:
if (bit_width_impl_type is None
or bit_width is None
or restrict_bit_width_type is None):
raise Exception("Bit width is not defined properly")
if bit_width_impl_type == BitWidthImplType.CONST:
tensor_clamp_impl = TensorClampSte()
bit_width_impl = BitWidthConst(bit_width, restrict_bit_width_type)
elif bit_width_impl_type == BitWidthImplType.PARAMETER:
tensor_clamp_impl = TensorClamp()
bit_width_impl = BitWidthParameter(bit_width_init=bit_width,
restrict_bit_width_type=restrict_bit_width_type,
min_overall_bit_width=min_overall_bit_width,
max_overall_bit_width=max_overall_bit_width,
override_pretrained=override_pretrained_bit_width)
else:
raise Exception("Bit width type {} not supported for weight quantization."
.format(str(bit_width_impl_type)))
else:
tensor_clamp_impl = TensorClamp()
bit_width_impl = bit_width_impl_override
float_to_int_impl = RestrictValue(restrict_value_type=RestrictValueType.INT,
float_to_int_impl_type=FloatToIntImplType.ROUND,
min_val=None)
int_scaling_impl = IntScaling(narrow_range,
signed=True,
restrict_scaling_type=restrict_scaling_type)
tensor_quant = RescalingIntQuant(narrow_range=narrow_range,
signed=True,
scaling_impl=scaling_impl,
int_scaling_impl=int_scaling_impl,
tensor_clamp_impl=tensor_clamp_impl,
msb_clamp_bit_width_impl=bit_width_impl,
float_to_int_impl=float_to_int_impl,
runtime=False)
else:
raise Exception('Unsupported weight quantization: {} bit width, {} quantization.'
.format(bit_width, str(quant_type)))
return tensor_quant
def __init__(self,
activation_impl: Module,
bit_width: int,
signed: bool,
narrow_range: bool,
min_val: float,
max_val: float,
quant_type: QuantType,
float_to_int_impl_type: FloatToIntImplType,
scaling_override: Optional[Module],
scaling_impl_type: ScalingImplType,
scaling_per_channel: bool,
scaling_min_val: Optional[float],
scaling_stats_sigma: Optional[float],
scaling_stats_op: Optional[StatsOp],
scaling_stats_buffer_momentum: Optional[float],
scaling_stats_permute_dims: Optional[Tuple],
per_channel_broadcastable_shape: Optional[Tuple[int, ...]],
min_overall_bit_width: Optional[int],
max_overall_bit_width: Optional[int],
bit_width_impl_override: Module,
bit_width_impl_type: BitWidthImplType,
restrict_bit_width_type: RestrictValueType,
restrict_scaling_type: RestrictValueType,
override_pretrained_bit_width: bool):
super(ActivationQuantProxy, self).__init__()
if not signed and min_val != 0.0:
raise Exception("Min val has to be 0.0 when quantization is unsigned.")
if scaling_per_channel and per_channel_broadcastable_shape is None:
raise Exception("Per channel scaling requires to specify number of channels.")
if quant_type == QuantType.FP:
tensor_quant = IdentityQuant()
else:
if scaling_impl_type != ScalingImplType.OVERRIDE and scaling_override is not None:
raise Exception("Overriding scaling requires to set ScalingImplType to OVERRIDE explicitly.")
if scaling_impl_type == ScalingImplType.OVERRIDE and scaling_override is None:
raise Exception("Overriding scaling requires to pass a scaling impl module.")
if scaling_per_channel:
scaling_shape = per_channel_broadcastable_shape
else:
scaling_shape = SCALING_SCALAR_SHAPE
if scaling_impl_type == ScalingImplType.OVERRIDE and scaling_override is not None:
scaling_impl = scaling_override
runtime = False
elif scaling_impl_type == ScalingImplType.CONST or scaling_impl_type == ScalingImplType.PARAMETER:
scaling_init = RescalingIntQuant.scaling_init_from_min_max(min_val, max_val)
scaling_impl = StandaloneScaling(is_parameter=scaling_impl_type == ScalingImplType.PARAMETER,
parameter_shape=scaling_shape,
restrict_scaling_type=restrict_scaling_type,
scaling_init=scaling_init,
scaling_min_val=scaling_min_val)
runtime = False
elif scaling_impl_type == ScalingImplType.STATS or scaling_impl_type == ScalingImplType.AFFINE_STATS:
if scaling_per_channel and not scaling_stats_op == StatsOp.MAX_AVE:
scaling_stats_input_view_shape_impl = StatsInputViewShapeImpl.OVER_OUTPUT_CHANNELS
scaling_stats_reduce_dim = 1
elif scaling_per_channel and scaling_stats_op == StatsOp.MAX_AVE:
raise Exception("Can't do per channel scaling with MAX AVE statistics.")
elif not scaling_per_channel and scaling_stats_op == StatsOp.MAX_AVE:
scaling_stats_input_view_shape_impl = StatsInputViewShapeImpl.OVER_OUTPUT_CHANNELS
scaling_stats_reduce_dim = 1
else: # not scaling_per_channel
scaling_stats_input_view_shape_impl = StatsInputViewShapeImpl.OVER_TENSOR
scaling_stats_reduce_dim = None
scaling_stats_permute_dims = None
stats_buffer_init = RescalingIntQuant.scaling_init_from_min_max(min_val, max_val).item()
scaling_impl = RuntimeStatsScaling(stats_op=scaling_stats_op,
restrict_scaling_type=restrict_scaling_type,
stats_input_view_shape_impl=scaling_stats_input_view_shape_impl,
stats_output_shape=scaling_shape,
sigma=scaling_stats_sigma,
scaling_min_val=scaling_min_val,
stats_reduce_dim=scaling_stats_reduce_dim,
stats_buffer_momentum=scaling_stats_buffer_momentum,
stats_buffer_init=stats_buffer_init,
stats_permute_dims=scaling_stats_permute_dims,
affine=scaling_impl_type == ScalingImplType.AFFINE_STATS)
runtime = True
else:
raise Exception("Scaling type {} not supported for int runtime quantization"
.format(str(scaling_impl_type)))
if quant_type == QuantType.BINARY:
if not signed:
raise Exception("Binary activation supports only signed activations")
tensor_quant = ClampedBinaryQuant(scaling_impl=scaling_impl)
elif quant_type == QuantType.INT:
if bit_width_impl_override is None:
if bit_width_impl_type is None or bit_width is None or restrict_bit_width_type is None:
raise Exception("Bit width is not defined properly")
if bit_width_impl_type == BitWidthImplType.CONST:
tensor_clamp_impl = TensorClamp() # If it's const, don't pass gradients to clipped values
msb_clamp_bit_width_impl = BitWidthConst(bit_width, restrict_bit_width_type)
elif bit_width_impl_type == BitWidthImplType.PARAMETER:
tensor_clamp_impl = TensorClamp() # if it's learned, I pass gradients to the bit width
msb_clamp_bit_width_impl = BitWidthParameter(bit_width,
min_overall_bit_width,
max_overall_bit_width,
restrict_bit_width_type,
override_pretrained_bit_width)
else:
raise Exception("Bit width type {} not supported for weight quantization"
.format(str(bit_width_impl_type)))
else:
msb_clamp_bit_width_impl = bit_width_impl_override
tensor_clamp_impl = TensorClamp() # if there is an override, it's learned
float_to_int_impl = RestrictValue(restrict_value_type=RestrictValueType.INT,
float_to_int_impl_type=float_to_int_impl_type,
min_val=None)
int_scaling_impl = IntScaling(narrow_range,
signed=signed,
restrict_scaling_type=restrict_scaling_type)
tensor_quant = RescalingIntQuant(signed=signed,
narrow_range=narrow_range,
scaling_impl=scaling_impl,
int_scaling_impl=int_scaling_impl,
tensor_clamp_impl=tensor_clamp_impl,
msb_clamp_bit_width_impl=msb_clamp_bit_width_impl,
float_to_int_impl=float_to_int_impl,
runtime=runtime)
else:
raise Exception("Quantization type {} not supported for activations.".format(quant_type))
self.fused_activation_quant_proxy = FusedActivationQuantProxy(activation_impl, tensor_quant)
self.scaling_impl_type = scaling_impl_type # needed to switch between different scaling modes