def test_bias_correction(self):
# Generating random numbers from a normal distribution for the weights and biases of the current and prev layer
np.random.seed(1)
shape = (2, 3, 2, 2)
# output 1
o1 = np.random.randn(*shape)
# output 2
o2 = np.random.randn(*shape)
biasCorrection = libpymo.BiasCorrection()
biasCorrection.storePreActivationOutput(o1)
biasCorrection.storePreActivationOutput(o1)
biasCorrection.storeQuantizedPreActivationOutput(o2)
biasCorrection.storeQuantizedPreActivationOutput(o2)
bias_tensor = libpymo.TensorParamBiasCorrection()
bias = np.array(np.random.randn(shape[1]))
bias_tensor.data = bias
biasCorrection.correctBias(bias_tensor)
bias_python = correct_bias(o1, o2, bias)
print(bias_tensor.data)
print(bias_python)
assert np.allclose(bias_tensor.data, bias_python)
def test_bias_update(self):
np.random.seed(1)
layer = nn.Conv2d(3, 10, 5)
bias_corr = libpymo.BiasCorrection()
bias_before = layer.bias.detach().cpu().numpy()
shape = (10, 10, 5, 5)
reference_output_batch = np.random.randn(*shape)
quantized_model_output_batch = np.random.randn(*shape)
bias_corr.storePreActivationOutput(reference_output_batch)
bias_corr.storeQuantizedPreActivationOutput(
quantized_model_output_batch)
bias_correction.call_empirical_mo_correct_bias(layer, bias_corr)
bias_after = layer.bias.detach().cpu().numpy()
# Assert bias has changed after running bias correction
self.assertFalse(np.allclose(bias_before, bias_after))
def correct_bias(model: torch.nn.Module,
quant_params: qsim.QuantParams,
num_quant_samples: int,
data_loader,
num_bias_correct_samples: int,
conv_bn_dict: Union[Dict[torch.nn.Module, ConvBnInfoType],
None] = None,
perform_only_empirical_bias_corr: bool = True,
layers_to_ignore: List[torch.nn.Module] = None):
"""
Corrects bias for each Conv layer of model (unless ignored). A combination of Analytical and Empirical Bias
Correction is used i.e. all the layers which can be corrected using Analytical Bias Correction are corrected
using Analytical Bias Correction and remaining layers are corrected using Empirical method.
Returns an in-place corrected floating point model
:param model: Model to be corrected
:param quant_params: Named tuple for quantization simulation for bias correction
:param num_quant_samples: number of samples of images to pass through quantization sim for bias correction.
:param data_loader: data loader for the model
:param num_bias_correct_samples: number of samples for Bias correction
:param conv_bn_dict: Dict of conv and bn with information related to activation. If None, the function calc it
:param perform_only_empirical_bias_corr: Default True. If true will perform only empirical Bias Corr for all layers
irrespective of the fact that layer is eligible for Analytical Bias Corr.
:param layers_to_ignore: list of layer names for which we need to skip bias correction.
"""
if layers_to_ignore is None:
layers_to_ignore = []
# Find batch size and shape of input tensor
batch_size, input_shape = utils.get_input_shape_batch_size(data_loader)
# Rounding up number of samples to batch size
n_batches_bias_correction = int(
np.ceil(num_bias_correct_samples / batch_size))
n_batches_quantization = int(np.ceil(num_quant_samples / batch_size))
data_loader_n_samples_bias_corr = utils.IterFirstX(
data_loader, n_batches_bias_correction)
data_loader_n_samples_quant = utils.IterFirstX(data_loader,
n_batches_quantization)
# TODO: Remove wrapper function
# Create a wrapping function for data loader for quantization
def pass_data_through_model(model,
early_stopping_iterations=None,
use_cuda=False):
# pylint: disable=unused-argument
# forward pass for given number of batches for model
for (images_in_one_batch, _) in data_loader_n_samples_quant:
forward_pass(model, images_in_one_batch)
ordered_conv_linear_nodes = get_ordered_lists_of_conv_fc(
model, input_shape)
if conv_bn_dict is None:
conv_bn_dict = find_all_conv_bn_with_activation(model, input_shape)
# Create a copy of the model as reference model
model_copy = copy.deepcopy(model)
# Add bias for all the layers whose bias is None
for name, module in ordered_conv_linear_nodes:
if module.bias is None:
if isinstance(module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):
output_size = module.out_channels
elif isinstance(module, torch.nn.Linear):
output_size = module.out_features
module.bias = torch.nn.Parameter(torch.zeros(output_size))
module.bias.data = module.bias.data.to(device=module.weight.device)
# Quantize full model
dummy_tensors = utils.create_rand_tensors_given_shapes(input_shape)
dummy_tensors = [
tensor.to(utils.get_device(model)) for tensor in dummy_tensors
]
q = qsim.QuantizationSimModel(model=model,
quant_scheme=quant_params.quant_scheme,
rounding_mode=quant_params.round_mode,
default_output_bw=quant_params.act_bw,
default_param_bw=quant_params.weight_bw,
in_place=True,
dummy_input=dummy_tensors,
config_file=quant_params.config_file)
# make sure model got updated in-place before we use it for bc updates
assert (q.model is model)
# updates to skip_output_activation and layers_to_ignore
for name, module in model.named_modules():
# Skip all layer's output quantization
if isinstance(module, QcQuantizeWrapper):
module.output_quantizers[0].enabled = False
q.compute_encodings(pass_data_through_model, None)
# For first conv layer, perform analytical bc if perform_only_empirical_bias_corr is set to False
# and layer is not marked to be ignored during bc.
if not perform_only_empirical_bias_corr:
module_name, module = ordered_conv_linear_nodes[0]
if module not in layers_to_ignore:
logger.info('Correcting layer %s using Analytical Bias Correction',
module_name)
quantize_layer = utils.get_layer_by_name(model, module_name)
call_analytical_mo_correct_bias(quantize_layer, None, None)
logger.info('Corrected bias for the layer')
ordered_conv_linear_nodes.pop(0)
for module_name, module in ordered_conv_linear_nodes:
# Ignore all layers which are skipped by user
if module in layers_to_ignore:
continue
else:
# make sure module is in the model used by qsim.
assert (module in list(q.model.modules()))
# Analytical Bias Correction is only done for Conv layers
reference_layer = utils.get_layer_by_name(model_copy, module_name)
quantize_layer = utils.get_layer_by_name(model, module_name)
if module in conv_bn_dict.keys():
bn_layer_info = conv_bn_dict[module]
if perform_only_empirical_bias_corr or bn_layer_info is None or bn_layer_info.input_bn is None:
logger.info(
'Correcting layer %s using Empirical Bias Correction',
module_name)
bias_correction = libpymo.BiasCorrection()
# Get output from quantized model and reference model
for images_in_one_batch, _ in data_loader_n_samples_bias_corr:
reference_output_batch = get_output_data(
reference_layer, model_copy, images_in_one_batch)
quantized_model_output_batch = get_output_data(
quantize_layer, model, images_in_one_batch)
if isinstance(reference_layer, torch.nn.Linear):
extended_shape = np.concatenate(
(reference_output_batch.shape, np.array([1,
1])))
reference_output_batch = reference_output_batch.reshape(
extended_shape)
quantized_model_output_batch = quantized_model_output_batch.reshape(
extended_shape)
bias_correction.storePreActivationOutput(
reference_output_batch)
bias_correction.storeQuantizedPreActivationOutput(
quantized_model_output_batch)
call_empirical_mo_correct_bias(module, bias_correction)
else:
logger.info(
'Correcting layer %s using Analytical Bias Correction',
module_name)
call_analytical_mo_correct_bias(
quantize_layer, bn_layer_info.input_bn,
bn_layer_info.in_activation_type)
logger.info('Corrected bias for the layer')
SaveUtils.remove_quantization_wrappers(model)
logger.info('Completed bias correction')
def my_correct_bias(model: torch.nn.Module,
quant_params,
num_quant_samples: int,
data_loader,
num_bias_correct_samples: int,
conv_bn_dict=None,
perform_only_empirical_bias_corr: bool = True,
layers_to_ignore=None,
quantizer_modifications=None):
if layers_to_ignore is None:
layers_to_ignore = []
# Find batch size and shape of input tensor
batch_size, input_shape = aimet_torch.utils.get_input_shape_batch_size(
data_loader)
# Rounding up number of samples to batch size
n_batches_bias_correction = int(
np.ceil(num_bias_correct_samples / batch_size))
n_batches_quantization = int(np.ceil(num_quant_samples / batch_size))
data_loader_n_samples_bias_corr = aimet_torch.utils.IterFirstX(
data_loader, n_batches_bias_correction)
data_loader_n_samples_quant = aimet_torch.utils.IterFirstX(
data_loader, n_batches_quantization)
# TODO: Remove wrapper function
# Create a wrapping function for data loader for quantization
def pass_data_through_model(model,
early_stopping_iterations=None,
use_cuda=False):
# pylint: disable=unused-argument
# forward pass for given number of batches for model
for (images_in_one_batch, _) in data_loader_n_samples_quant:
aimet_torch.bias_correction.forward_pass(model,
images_in_one_batch)
ordered_conv_linear_nodes = aimet_torch.utils.get_ordered_lists_of_conv_fc(
model, input_shape)
if conv_bn_dict is None:
conv_bn_dict = aimet_torch.bias_correction.find_all_conv_bn_with_activation(
model, input_shape)
# Create a copy of the model as reference model
model_copy = copy.deepcopy(model)
# Add bias for all the layers whose bias is None
for name, module in ordered_conv_linear_nodes:
if module.bias is None:
if isinstance(module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)):
output_size = module.out_channels
elif isinstance(module, torch.nn.Linear):
output_size = module.out_features
module.bias = torch.nn.Parameter(torch.zeros(output_size))
module.bias.data = module.bias.data.to(device=module.weight.device)
# Quantize full model
dummy_tensors = aimet_torch.utils.create_rand_tensors_given_shapes(
input_shape)
dummy_tensors = [
tensor.to(aimet_torch.utils.get_device(model))
for tensor in dummy_tensors
]
q = aimet_torch.quantsim.QuantizationSimModel(
model=model,
quant_scheme=quant_params.quant_scheme,
rounding_mode=quant_params.round_mode,
default_output_bw=quant_params.act_bw,
default_param_bw=quant_params.weight_bw,
in_place=True,
dummy_input=dummy_tensors,
config_file=quant_params.config_file)
# make sure model got updated in-place before we use it for bc updates
assert (q.model is model)
if quantizer_modifications is not None:
quantizer_modifications(q)
# updates to skip_output_activation and layers_to_ignore
for name, module in model.named_modules():
# Skip all layer's output quantization
if isinstance(module, QcQuantizeWrapper):
module.output_quantizers[0].enabled = False
q.compute_encodings(pass_data_through_model, None)
# For first conv layer, perform analytical bc if perform_only_empirical_bias_corr is set to False
# and layer is not marked to be ignored during bc.
if not perform_only_empirical_bias_corr:
module_name, module = ordered_conv_linear_nodes[0]
if module not in layers_to_ignore:
aimet_torch.bias_correction.logger.info(
'Correcting layer %s using Analytical Bias Correction',
module_name)
quantize_layer = aimet_torch.utils.get_layer_by_name(
model, module_name)
aimet_torch.bias_correction.call_analytical_mo_correct_bias(
quantize_layer, None, None)
aimet_torch.bias_correction.logger.info(
'Corrected bias for the layer')
ordered_conv_linear_nodes.pop(0)
for module_name, module in ordered_conv_linear_nodes:
# Ignore all layers which are skipped by user
if module in layers_to_ignore or module_name in layers_to_ignore:
continue
else:
# make sure module is in the model used by qsim.
assert (module in list(q.model.modules()))
# Analytical Bias Correction is only done for Conv layers
reference_layer = aimet_torch.utils.get_layer_by_name(
model_copy, module_name)
quantize_layer = aimet_torch.utils.get_layer_by_name(
model, module_name)
if module in conv_bn_dict.keys():
bn_layer_info = conv_bn_dict[module]
if perform_only_empirical_bias_corr or bn_layer_info is None or bn_layer_info.input_bn is None:
aimet_torch.bias_correction.logger.info(
'Correcting layer %s using Empirical Bias Correction',
module_name)
bias_correction = libpymo.BiasCorrection()
# Get output from quantized model and reference model
for images_in_one_batch, _ in data_loader_n_samples_bias_corr:
reference_output_batch = aimet_torch.bias_correction.get_output_data(
reference_layer, model_copy, images_in_one_batch)
quantized_model_output_batch = aimet_torch.bias_correction.get_output_data(
quantize_layer, model, images_in_one_batch)
if isinstance(reference_layer, torch.nn.Linear):
extended_shape = np.concatenate(
(reference_output_batch.shape, np.array([1,
1])))
reference_output_batch = reference_output_batch.reshape(
extended_shape)
quantized_model_output_batch = quantized_model_output_batch.reshape(
extended_shape)
bias_correction.storePreActivationOutput(
reference_output_batch)
bias_correction.storeQuantizedPreActivationOutput(
quantized_model_output_batch)
aimet_torch.bias_correction.call_empirical_mo_correct_bias(
module, bias_correction)
else:
aimet_torch.bias_correction.logger.info(
'Correcting layer %s using Analytical Bias Correction',
module_name)
aimet_torch.bias_correction.call_analytical_mo_correct_bias(
quantize_layer, bn_layer_info.input_bn,
bn_layer_info.in_activation_type)
aimet_torch.bias_correction.logger.info(
'Corrected bias for the layer')
aimet_torch.save_utils.SaveUtils.remove_quantization_wrappers(model)
aimet_torch.bias_correction.logger.info('Completed bias correction')
def bias_correction_per_layer(
reference_model: tf.compat.v1.Session,
corrected_model: tf.compat.v1.Session,
bias_correct_params: BiasCorrectionParams,
layer_name_to_be_corrected: str, quant_params: QuantParams,
data_set: tf.data.Dataset) -> tf.compat.v1.Session:
"""
Helper function to perform empirical bias correction per layer.
:param reference_model: active tensorflow session for reference model
:param corrected_model: active tensorflow session for corrected model
:param bias_correct_params: bias correction params
:param layer_name_to_be_corrected: name of layer on which bias correction is to be performed
:param quant_params: Quantization specific params from user
:return: None, updates corrected model in-place.
"""
# Quantize model
quantize_model = BiasCorrection._get_quantized_model(
corrected_model, quant_params, bias_correct_params.input_op_names,
bias_correct_params.output_op_names,
bias_correct_params.num_quant_samples,
bias_correct_params.batch_size, data_set)
ref_layer = reference_model.graph.get_operation_by_name(
layer_name_to_be_corrected)
bias_correction = libpymo.BiasCorrection()
logger.info('Correcting layer %s', ref_layer.name)
n_batches_bias_correction = int(
np.ceil(bias_correct_params.num_bias_correct_samples /
bias_correct_params.batch_size))
reduced_dataset_iter = iter_first_x(data_set,
n_batches_bias_correction)
for batch_input in reduced_dataset_iter:
# reference model without corrected nodes
reference_output_batch = BiasCorrection._get_output_data(
reference_model, bias_correct_params.input_op_names,
ref_layer.name, batch_input)
quantized_model_output_batch = BiasCorrection._get_output_data(
quantize_model, bias_correct_params.input_op_names,
ref_layer.name, batch_input)
if ref_layer.type == 'MatMul':
extended_shape = np.concatenate(
(reference_output_batch.shape, np.array([1, 1])))
reference_output_batch = reference_output_batch.reshape(
extended_shape)
quantized_model_output_batch = quantized_model_output_batch.reshape(
extended_shape)
# we need to reshape from tensorflow shape NxHxWxC to NxCxHxW
bias_correction.storePreActivationOutput(
np.ascontiguousarray(reference_output_batch.transpose(0, 3, 1, 2)))
bias_correction.storeQuantizedPreActivationOutput(
np.ascontiguousarray(
quantized_model_output_batch.transpose(0, 3, 1, 2)))
bias_shape = None
# get shape for bias if the layer does not have bias
if BiasUtils.is_bias_none(ref_layer):
if ref_layer.type == 'MatMul':
bias_shape = reference_output_batch.shape[1]
elif ref_layer.type in ['Conv2D', 'DepthwiseConv2dNative']:
# for conv2d or depthwise conv2d
bias_shape = reference_output_batch.shape[3]
# bias is to be corrected in the corrected model graph
BiasCorrection._call_mo_correct_bias(corrected_model, ref_layer.name,
bias_correction, bias_shape)
logger.info('Completed empirical bias correction for layer %s',
ref_layer.name)