def _get_conv_linear_params(model, layer_to_be_corrected):
"""
Extract weights and bias of given conv/linear layer
:param model: tf.compat.v1.Session type
:param layer_to_be_corrected: conv/linear layer as tf.Operation
:return: bias, weight and quantized weights as TensorParamBiasCorrection types
"""
bias_tensor = libpymo.TensorParamBiasCorrection()
# get weight tensor
weight_tensor, _ = get_weight_tensor_with_shape(
model, layer_to_be_corrected)
if weight_tensor is None:
logger.error('Weight tensor extraction failed for layer {%s}',
layer_to_be_corrected.name)
# get bias tensor, at this point we have initialized model layers to have bias add param.
assert not BiasUtils.is_bias_none(layer_to_be_corrected)
bias_tensor.data = BiasUtils.get_bias_as_numpy_data(
model, layer_to_be_corrected)
bias_tensor.shape = BiasUtils.get_shape(layer_to_be_corrected)
return bias_tensor, weight_tensor
def test_bias_add_with_conv(self):
"""
Test bias add on conv op
:return:
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(32, 32, 3,), name="inputs")
# create a conv without bias param
conv_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv_op)
# pylint: disable=no-member
_ = tf.nn.relu(bn_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
self.assertTrue(BiasUtils.is_bias_none(conv_op))
# new_sess = BiasUtils.initialize_model_with_bias(sess)
shape = BiasUtils.get_shape(conv_op)
numpy_data = np.random.rand(shape[0])
BiasUtils.update_bias_for_op(sess, conv_op, bias_as_numpy_array=numpy_data)
new_sess = save_and_load_graph('./temp_bn_fold', sess)
conv_op = new_sess.graph.get_operation_by_name('conv2d/Conv2D')
bias_as_numpy_data = BiasUtils.get_bias_as_numpy_data(new_sess, conv_op)
assert(not BiasUtils.is_bias_none(conv_op))
new_sess.close()
def scale_cls_set_with_conv_layers(model: tf.compat.v1.Session, cls_set: Tuple[tf.Operation, tf.Operation]) -> np.ndarray:
"""
API to invoke equalize layer params (update for weights and bias is in place)
:param model: active tf.compat.v1.Session
:param cls_set: Consecutive Conv layers Tuple whose weights and biases need to be equalized
:return: Scaling factor S_12 for each conv layer pair: numpy array
"""
with model.graph.as_default():
for module in cls_set:
if module.type not in ['Conv2D', 'DepthwiseConv2dNative']:
raise ValueError("Only conv layers are supported for cross layer equalization")
# Create structs for holding layer weights and bias parameters
prev_layer_params = libpymo.EqualizationParams()
curr_layer_params = libpymo.EqualizationParams()
# send as [Noc, Nic, kh, kw], TF format is [kh, kw, Nic, Noc]
prev_layer_params.weight = WeightTensorUtils.get_tensor_as_numpy_data(model, cls_set[0]). \
transpose((3, 2, 0, 1)).reshape(-1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_set[0])
prev_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0], weight_shape[1]]
prev_layer_params.isBiasNone = BiasUtils.is_bias_none(cls_set[0])
# send as [Noc, Nic, kh, kw], TF format is [kh, kw, Nic, Noc]
curr_layer_params.weight = WeightTensorUtils.get_tensor_as_numpy_data(model, cls_set[1]). \
transpose((3, 2, 0, 1)).reshape(-1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_set[1])
curr_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0], weight_shape[1]]
if not BiasUtils.is_bias_none(cls_set[0]):
prev_layer_params.bias = BiasUtils.get_bias_as_numpy_data(model, cls_set[0]).reshape(-1)
else:
prev_layer_params.isBiasNone = True
scaling_factor = libpymo.scaleLayerParams(prev_layer_params, curr_layer_params)
# convert received formats back to TF
# TF format is [kh, kw, Nic, Noc]
numpy_weight_reshaped = np.reshape(prev_layer_params.weight, prev_layer_params.weightShape). \
transpose((2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(model, cls_set[0], numpy_weight_reshaped)
numpy_weight_reshaped = np.reshape(curr_layer_params.weight, curr_layer_params.weightShape). \
transpose((2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(model, cls_set[1], numpy_weight_reshaped)
if not BiasUtils.is_bias_none(cls_set[0]):
numpy_bias_reshaped = np.reshape(prev_layer_params.bias, BiasUtils.get_shape(cls_set[0]))
BiasUtils.update_bias_for_op(model, cls_set[0], numpy_bias_reshaped)
return scaling_factor
def _get_bias_tensor(sess: tf.compat.v1.Session,
conv: tf.Operation) -> libpymo.TensorParams():
"""
Get bias tensor in given conv op.
Packs bias in the format required for BN fold
(libpymo.TensorParams()).
:param sess: current session
:param conv: conv op
:return: return bias param in libpymo.TensorParams() format.
"""
# Bias tensor
bias_tensor = libpymo.TensorParams()
with sess.graph.as_default():
if not BiasUtils.is_bias_none(conv):
bias_tensor.shape = BiasUtils.get_shape(conv)
bias_tensor.data = BiasUtils.get_bias_as_numpy_data(sess, conv)
return bias_tensor
def test_update_to_bias_with_load_var(self):
"""
tests update to bias param of conv op using tf variable load api
:return:
"""
# create conv op
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(32, 32, 3,))
conv_op = tf.keras.layers.Conv2D(32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 2))(inputs)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv_op)
_ = tf.nn.relu(bn_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
original_bias = BiasUtils.get_bias_as_numpy_data(sess, conv_op)
# add dummy weight tensor data
np.random.seed(0)
b_shape = BiasUtils.get_shape(conv_op)
numpy_data = np.random.rand(b_shape[0])
# send in numpy data to overwrite previous value
BiasUtils.update_bias_for_op(sess, conv_op, numpy_data)
updated_bias = BiasUtils.get_bias_as_numpy_data(sess, conv_op)
# validate they are not the same
self.assertFalse(np.allclose(original_bias, updated_bias))
self.assertTrue(np.allclose(numpy_data, updated_bias))
sess.close()
def bias_fold(sess: tf.compat.v1.Session, folded_pairs: List[Tuple[tf.Operation, tf.Operation]],
cls_set_info_list: List[ClsSetInfo]) -> tf.compat.v1.Session:
"""
Folds bias values greater than 3 * sigma to next layer's bias
:param sess: Current session
:param folded_pairs: Key: Conv/Linear layer Value: Corresponding folded BN layer
:param cls_set_info_list: List of info elements for each cls set
:return: updated session after graph updates from hbf
"""
with sess.graph.as_default():
# refresh the references saved during bn fold and cls.
cls_set_info_list, bn_layers = HighBiasFold._refresh_layer_set_info_before_hbf(sess, folded_pairs,
cls_set_info_list)
if not bn_layers:
logger.error('High Bias folding is not supported for models without BatchNorm Layers')
return sess
for cls_set_info in cls_set_info_list:
for cls_pair_info in cls_set_info.cls_pair_info_list:
# check if we have a corresponding bn layer
if cls_pair_info.layer1.name in bn_layers.keys():
# check if bias present in given conv2D(s)
if BiasUtils.is_bias_none(cls_pair_info.layer1) or BiasUtils.is_bias_none(cls_pair_info.layer2):
continue
prev_layer_params = libpymo.LayerParams()
curr_layer_params = libpymo.LayerParams()
scaling_parameter = cls_pair_info.scale_factor
prev_layer_bn_params =\
HighBiasFold.get_bn_params_for_bias_fold(sess,
bn_layers[cls_pair_info.layer1.name],
scaling_parameter)
prev_layer_params.activationIsRelu = cls_pair_info.relu_activation_between_layers
prev_layer_params.bias =\
BiasUtils.get_bias_as_numpy_data(sess, cls_pair_info.layer1).reshape(-1)
prev_bias_shape = BiasUtils.get_shape(cls_pair_info.layer1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_pair_info.layer1)
prev_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0],
weight_shape[1]]
curr_layer_params.bias =\
BiasUtils.get_bias_as_numpy_data(sess, cls_pair_info.layer2).reshape(-1)
curr_bias_shape = BiasUtils.get_shape(cls_pair_info.layer2)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_pair_info.layer2)
# Handle depthwise layer case
# for a depthwise layer num outputs is set to 1 in TF
# send as [Nic, Noc, kh, kw], TF format is [kh, kw, Nic, Noc]
if cls_pair_info.layer2.type in ['DepthwiseConv2dNative']:
c_wt = WeightTensorUtils.get_tensor_as_numpy_data(
sess, cls_pair_info.layer2).transpose((2, 3, 0, 1))
curr_layer_params.weight = c_wt.reshape(-1)
curr_layer_params.weightShape = [weight_shape[2], weight_shape[3], weight_shape[0],
weight_shape[1]]
else:
# send as [Noc, Nic, kh, kw], TF format is [kh, kw, Nic, Noc]
c_wt = WeightTensorUtils.get_tensor_as_numpy_data(
sess, cls_pair_info.layer2).transpose((3, 2, 0, 1))
curr_layer_params.weight = c_wt.reshape(-1)
curr_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0],
weight_shape[1]]
libpymo.updateBias(prev_layer_params, curr_layer_params, prev_layer_bn_params)
BiasUtils.update_bias_for_op(sess, cls_pair_info.layer1, np.reshape(prev_layer_params.bias,
prev_bias_shape))
BiasUtils.update_bias_for_op(sess, cls_pair_info.layer2, np.reshape(curr_layer_params.bias,
curr_bias_shape))
else:
logger.info("skipping layer: {%s}", cls_pair_info.layer1.name)
# save and load the updated graph after high bias fold update
aftr_hbf_sess = save_and_load_graph('./temp_hbf', sess)
return aftr_hbf_sess
def scale_cls_set_with_depthwise_layers(model: tf.compat.v1.Session,
cls_set: Tuple[tf.Operation,
tf.Operation,
tf.Operation]) -> [np.ndarray, np.ndarray]:
"""
API to invoke equalize layer params for depth wise separable layers(update for weights and bias is in place)
:param model: active tf.compat.v1.Session
:param cls_set: Consecutive Conv layers whose weights and biases need to be equalized.
Second Conv layer is a depth-wise conv and third conv layer is point-wise conv
:return: Scaling factors S_12 and S_23 : numpy arrays
"""
# make sure you define the session and graph scope before making any graph updates.
with model.graph.as_default():
for module in cls_set:
if module.type not in ['Conv2D', 'DepthwiseConv2dNative']:
raise ValueError("Only conv layers are supported for cross layer equalization")
# Create structs for holding layer weights and bias parameters
prev_layer_params = libpymo.EqualizationParams()
curr_layer_params = libpymo.EqualizationParams()
next_layer_params = libpymo.EqualizationParams()
# send as [Noc, Nic, kh, kw], TF format is [kh, kw, Nic, Noc]
prev_layer_params.weight = WeightTensorUtils.get_tensor_as_numpy_data(model, cls_set[0]). \
transpose((3, 2, 0, 1)).reshape(-1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_set[0])
prev_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0], weight_shape[1]]
prev_layer_params.isBiasNone = BiasUtils.is_bias_none(cls_set[0])
# depthwise layer outputs is set to 1 in TF
# send as [Nic, Noc, kh, kw], TF format is [kh, kw, Nic, Noc]
curr_layer_params.weight = WeightTensorUtils.get_tensor_as_numpy_data(model, cls_set[1]). \
transpose((2, 3, 0, 1)).reshape(-1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_set[1])
# depthwise layer outputs is set to 1 in TF
# send as [Nic, Noc, kh, kw], TF format is [kh, kw, Nic, Noc]
curr_layer_params.weightShape = [weight_shape[2], weight_shape[3], weight_shape[0], weight_shape[1]]
curr_layer_params.isBiasNone = BiasUtils.is_bias_none(cls_set[1])
# send as [Noc, Nic, kh, kw] , TF format is [kh, kw, Nic, Noc]
next_layer_params.weight = WeightTensorUtils.get_tensor_as_numpy_data(model, cls_set[2]). \
transpose((3, 2, 0, 1)).reshape(-1)
weight_shape = WeightTensorUtils.get_tensor_shape(cls_set[2])
next_layer_params.weightShape = [weight_shape[3], weight_shape[2], weight_shape[0], weight_shape[1]]
if not BiasUtils.is_bias_none(cls_set[0]):
prev_layer_params.bias = BiasUtils.get_bias_as_numpy_data(model, cls_set[0]).reshape(-1)
else:
prev_layer_params.isBiasNone = True
if not BiasUtils.is_bias_none(cls_set[1]):
curr_layer_params.bias = BiasUtils.get_bias_as_numpy_data(model, cls_set[1]).reshape(-1)
else:
curr_layer_params.isBiasNone = True
scaling_params = libpymo.scaleDepthWiseSeparableLayer(prev_layer_params, curr_layer_params,
next_layer_params)
# convert received formats back to TF
# TF format is [kh, kw, Nic, Noc]
numpy_weight_reshaped_0 = np.reshape(prev_layer_params.weight, prev_layer_params.weightShape). \
transpose((2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(model, cls_set[0], numpy_weight_reshaped_0)
# depthwise layer
numpy_weight_reshaped_1 = np.reshape(curr_layer_params.weight, curr_layer_params.weightShape). \
transpose((2, 3, 0, 1))
WeightTensorUtils.update_tensor_for_op(model, cls_set[1], numpy_weight_reshaped_1)
numpy_weight_reshaped_2 = np.reshape(next_layer_params.weight, next_layer_params.weightShape). \
transpose((2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(model, cls_set[2], numpy_weight_reshaped_2)
if not BiasUtils.is_bias_none(cls_set[0]):
numpy_bias_reshaped = np.reshape(prev_layer_params.bias, BiasUtils.get_shape(cls_set[0]))
BiasUtils.update_bias_for_op(model, cls_set[0], numpy_bias_reshaped)
if not BiasUtils.is_bias_none(cls_set[1]):
numpy_bias_reshaped = np.reshape(curr_layer_params.bias, BiasUtils.get_shape(cls_set[1]))
BiasUtils.update_bias_for_op(model, cls_set[1], numpy_bias_reshaped)
return scaling_params.scalingMatrix12, scaling_params.scalingMatrix23
def test_bias_correction_single_layer(self):
"""
Test bias correction for a single layer api
"""
tf.compat.v1.reset_default_graph()
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
# create a custom model
inputs = tf.keras.Input(shape=(
32,
16,
3,
))
conv_op = tf.keras.layers.Conv2D(16, (3, 3))(inputs)
relu_1 = tf.nn.relu(conv_op)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_1)
relu_2 = tf.nn.relu(conv2_op)
# global initializer
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(config=config,
graph=tf.compat.v1.get_default_graph())
sess.run(init)
# populate conv with dummy bias and weights
np.random.seed(0)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
w_numpy_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2],
w_shape[3])
b_shape = BiasUtils.get_shape(conv_op)
b_numpy_data = np.random.rand(b_shape[0])
WeightTensorUtils.update_tensor_for_op(sess, conv_op, w_numpy_data)
BiasUtils.update_bias_for_op(sess, conv_op, b_numpy_data)
# save and load the updated graph after high bias fold update
n_sess = aimet_tensorflow.utils.graph_saver.save_and_load_graph(
'./test_update', sess)
output_op = n_sess.graph.get_operation_by_name('Relu_1')
conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
bias_data = BiasUtils.get_bias_as_numpy_data(n_sess, conv_op)
input_op_name = conv_op.inputs[0].op.name
bias_corr_input = BiasCorrectionParams(
batch_size=1,
num_quant_samples=10,
num_bias_correct_samples=10,
input_op_names=[input_op_name],
output_op_names=[output_op.name])
quant_params = QuantParams(use_cuda=False)
np.random.seed(0)
shape = conv_op.inputs[0].shape
dataset = np.random.rand(1, shape[1], shape[2], shape[3])
with unittest.mock.patch(
'aimet_tensorflow.bias_correction.iter_first_x'
) as iter_first_x:
iter_first_x.return_value = [dataset]
BiasCorrection.bias_correction_per_layer(
reference_model=n_sess,
corrected_model=n_sess,
bias_correct_params=bias_corr_input,
layer_name_to_be_corrected=conv_op.name,
quant_params=quant_params,
data_set=dataset)
conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
bias_data_updated = BiasUtils.get_bias_as_numpy_data(
n_sess, conv_op)
# needs a validation
self.assertFalse(np.allclose(bias_data, bias_data_updated, atol=1e-4))
print('Test completed')
sess.close()
n_sess.close()
def test_bias_correction_model_tf(self):
"""
Test bias correction for custom model
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
32,
32,
3,
))
conv_op = tf.keras.layers.Conv2D(32, (3, 3))(inputs)
relu_1 = tf.nn.relu(conv_op)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(relu_1)
relu_2 = tf.nn.relu(conv2_op)
conv3_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_2)
relu_3 = tf.nn.relu(conv3_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
# updating random bias and weight for one conv
np.random.seed(0)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
b_shape = BiasUtils.get_shape(conv_op)
numpy_data = np.random.rand(b_shape[0])
BiasUtils.update_bias_for_op(sess, conv_op, numpy_data)
w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
w_numpy_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2],
w_shape[3])
WeightTensorUtils.update_tensor_for_op(sess, conv_op, w_numpy_data)
# save and load the updated graph after high bias fold update
n_sess = save_and_load_graph('./test_update', sess)
output_op = n_sess.graph.get_operation_by_name('Relu_1')
conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
bias_data = BiasUtils.get_bias_as_numpy_data(n_sess, conv_op)
input_op_name = conv_op.inputs[0].op.name
output_op = n_sess.graph.get_operation_by_name('Relu_2')
input_op_names = [input_op_name]
output_op_names = [output_op.name]
batch_size = 1
num_samples = 10
np.random.seed(0)
shape = conv_op.inputs[0].shape
dataset = np.random.rand(10, 1, shape[1], shape[2], shape[3])
dataset = tf.convert_to_tensor(dataset)
dataset = tf.data.Dataset.from_tensor_slices(dataset)
quant_params = QuantParams(quant_mode='tf', use_cuda=False)
bias_correction_params = BiasCorrectionParams(
batch_size=batch_size,
num_quant_samples=num_samples,
num_bias_correct_samples=num_samples,
input_op_names=input_op_names,
output_op_names=output_op_names)
conv_op = sess.graph.get_operation_by_name('conv2d_1/Conv2D')
assert (BiasUtils.is_bias_none(conv_op))
new_sess = BiasCorrection.correct_bias(sess, bias_correction_params,
quant_params, dataset)
conv_op = new_sess.graph.get_operation_by_name('conv2d_1/Conv2D')
assert (not BiasUtils.is_bias_none(conv_op))
sess.close()
n_sess.close()
new_sess.close()