def test_bias_update_to_dense(self):
"""
test bias correction on matmul layer
:return:
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(32, 32, 3,))
x = tf.keras.layers.Flatten()(inputs)
dense = tf.keras.layers.Dense(2, use_bias=False, activation=tf.nn.softmax, name="single_residual")(x)
# pylint: disable=no-member
_ = tf.nn.relu(dense)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
sess.run(init)
dense_op = sess.graph.get_operation_by_name('single_residual/MatMul')
self.assertTrue(BiasUtils.is_bias_none(dense_op))
new_sess = BiasUtils.initialize_model_with_bias(sess, ['input_1'], ['Relu'])
dense_op = new_sess.graph.get_operation_by_name('single_residual/MatMul')
self.assertTrue(not BiasUtils.is_bias_none(dense_op))
new_sess.close()
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 split_module(layer: Layer, rank: int) -> (tf.Operation, tf.Operation):
"""
:param layer: Module to be split
:param rank: rank for splitting
:return: Two split modules
"""
h, v = SpatialSvdModuleSplitter.get_svd_matrices(layer, rank)
conv_a_stride, conv_b_stride = get_strides_for_split_conv_ops(
op=layer.module)
with layer.model.graph.as_default():
# Use last_last_index to get name of conv layer prior to '/Conv2D' suffix
last_slash_index = layer.module.name.rfind('/')
data_format = layer.module.get_attr('data_format').decode('utf-8')
if data_format == "NHWC":
data_format_channels = "channels_last"
else:
data_format_channels = "channels_first"
padding = layer.module.get_attr('padding').decode('utf-8')
# Conv weight indices follow 'HWIO'
conv_a_out = tf.keras.layers.Conv2D(filters=v.shape[3], kernel_size=(v.shape[0], v.shape[1]),
strides=conv_a_stride, name=layer.module.name[:last_slash_index] + '_a',
data_format=data_format_channels, padding=padding, use_bias=False)\
(layer.module.inputs[0])
# Update weights for conv_a
WeightTensorUtils.update_tensor_for_op(layer.model, conv_a_out.op,
v)
# get the succeeding bias tensor
bias_tensor = aimet_tensorflow.utils.common.get_succeeding_bias_tensor(
layer.module)
use_bias = bias_tensor is not None
conv_b_out = tf.keras.layers.Conv2D(filters=h.shape[3], kernel_size=(h.shape[0], h.shape[1]),
strides=conv_b_stride,
name=layer.module.name[:last_slash_index] + '_b',
data_format=data_format_channels, padding=padding, use_bias=use_bias)\
(conv_a_out)
if use_bias:
# Find and write bias value into newly created bias variable
bias_value = BiasUtils.get_bias_as_numpy_data(
layer.model, layer.module)
BiasUtils.update_bias_for_op(layer.model,
conv_b_out.op.inputs[0].op,
bias_value)
conv_b_out = conv_b_out.op.inputs[
0] # set conv_b_out to be output tensor of Conv2D op
# Update weights for conv_b
WeightTensorUtils.update_tensor_for_op(layer.model, conv_b_out.op,
h)
return conv_a_out.op, conv_b_out.op
def test_prune_model_tf_slim(self):
""" Punning a model with tf slim api """
# create tf.compat.v1.Session and initialize the weights and biases with zeros
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
# create session with graph
sess = tf.compat.v1.Session(graph=tf.Graph(), config=config)
with sess.graph.as_default():
# by default, model will be constructed in default graph
x = tf.compat.v1.placeholder(tf.float32, [1, 32, 32, 3])
_ = tf_slim_basic_model(x)
sess.run(tf.compat.v1.global_variables_initializer())
conn_graph_orig = ConnectedGraph(sess.graph, ['Placeholder'],
['tf_slim_model/Softmax'])
num_ops_orig = len(conn_graph_orig.get_all_ops())
# Create a layer database
orig_layer_db = LayerDatabase(model=sess,
input_shape=(1, 32, 32, 3),
working_dir=None)
conv1 = orig_layer_db.find_layer_by_name('Conv_1/Conv2D')
conv1_bias = BiasUtils.get_bias_as_numpy_data(orig_layer_db.model,
conv1.module)
layer_comp_ratio_list = [LayerCompRatioPair(conv1, Decimal(0.5))]
spatial_svd_pruner = SpatialSvdPruner()
comp_layer_db = spatial_svd_pruner.prune_model(orig_layer_db,
layer_comp_ratio_list,
CostMetric.mac,
trainer=None)
# Check that svd added these ops
_ = comp_layer_db.model.graph.get_operation_by_name('Conv_1_a/Conv2D')
_ = comp_layer_db.model.graph.get_operation_by_name('Conv_1_b/Conv2D')
conn_graph_new = ConnectedGraph(comp_layer_db.model.graph,
['Placeholder'],
['tf_slim_model/Softmax'])
num_ops_new = len(conn_graph_new.get_all_ops())
self.assertEqual(num_ops_orig + 1, num_ops_new)
bias_add_op = comp_layer_db.model.graph.get_operation_by_name(
'Conv_1_b/BiasAdd')
conv_1_b_op = comp_layer_db.model.graph.get_operation_by_name(
'Conv_1_b/Conv2D')
self.assertEqual(
conn_graph_new._module_identifier.get_op_info(bias_add_op),
conn_graph_new._module_identifier.get_op_info(conv_1_b_op))
self.assertTrue(
np.array_equal(
conv1_bias,
BiasUtils.get_bias_as_numpy_data(comp_layer_db.model,
conv_1_b_op)))
def create_conv2d_dense_type_params(my_op: Op):
""" Create products for conv2d, dense, depthwise conv2d, and similar """
tf_op = my_op.get_module()
weight_op = WeightTensorUtils.get_read_op(tf_op)
create_and_connect_product('kernel', weight_op.outputs[0].shape, my_op, weight_op.outputs[0])
if not BiasUtils.is_bias_none(tf_op):
bias_op = BiasUtils.get_bias_read_op(tf_op)
create_and_connect_product('bias', bias_op.outputs[0].shape, my_op, bias_op.outputs[0])
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 _update_layer_params(layer: Layer, new_weight: np.ndarray, new_bias: Union[np.ndarray, None]):
"""
update parameters (weights and bias) for given layer
:param layer: layer to be updated
:param new_weight: newer weights
:param new_bias: new bias
:return:
"""
assert layer.module.type == 'Conv2D'
with layer.model.graph.as_default():
# update existing weight tensor with new_weight in place
WeightTensorUtils.update_tensor_for_op(layer.model, op=layer.module, tensor_as_numpy_array=new_weight)
if new_bias is not None:
# update existing bias tensor with new_bias in place
BiasUtils.update_bias_for_op(layer.model, op=layer.module, bias_as_numpy_array=new_bias)
def test_equalize_with_custom_model_no_bias(self):
"""
Test equalize with a custom model with conv without bias param
"""
tf.compat.v1.reset_default_graph()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
with sess.as_default():
inputs = tf.keras.Input(shape=(
32,
32,
3,
))
conv_op = tf.keras.layers.Conv2D(32, (3, 3),
use_bias=False)(inputs)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv_op)
relu_1 = tf.nn.relu(bn_op)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3),
use_bias=False)(relu_1)
bn_op_2 = tf.keras.layers.BatchNormalization(fused=True)(
conv2_op, training=False)
relu_2 = tf.nn.relu(bn_op_2)
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
old_conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
self.assertTrue(BiasUtils.is_bias_none(old_conv_op))
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
new_sess = equalize_model(sess, conv_op.inputs[0].op.name,
'Relu_1')
new_conv_op = new_sess.graph.get_operation_by_name('conv2d/Conv2D')
bias = BiasUtils.get_bias_as_numpy_data(new_sess, new_conv_op)
self.assertFalse(BiasUtils.is_bias_none(new_conv_op))
sess.close()
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 test_bias_add_custom_model(self):
""" test update bias when no bias present """
tf.compat.v1.reset_default_graph()
tf.set_random_seed(0)
inputs = tf.keras.Input(shape=(
32,
32,
3,
))
conv_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs)
relu2 = tf.nn.relu(conv2_op)
add = tf.keras.layers.add([conv_op, relu2])
relu = tf.nn.relu(add)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
shape = WeightTensorUtils.get_tensor_shape(conv_op.op)
np.random.seed(0)
bias_data = np.random.rand(shape[3])
assert BiasUtils.is_bias_none(conv_op.op)
BiasUtils.update_bias_for_op(sess, conv_op.op, bias_data)
n_sess = aimet_tensorflow.utils.graph_saver.save_and_load_graph(
'./test_update', sess)
conv_op_updated = n_sess.graph.get_operation_by_name(conv_op.op.name)
assert not BiasUtils.is_bias_none(conv_op_updated)
updated_bias = BiasUtils.get_bias_as_numpy_data(
n_sess, conv_op_updated)
self.assertTrue(np.allclose(updated_bias, bias_data))
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 test_equalize_fold_forward(self):
"""
Test equalize on a model with a forward bn fold
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
32,
32,
3,
), name="inputs")
conv_op = tf.keras.layers.Conv2D(32, (3, 3))(inputs)
r_op = tf.nn.relu(conv_op)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(r_op)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(bn_op)
conv3_op = tf.keras.layers.Conv2D(32, (3, 3))(conv2_op)
_ = tf.nn.relu(conv3_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
sess.run(init)
old_conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
conv_bias_data_before_fold = BiasUtils.get_bias_as_numpy_data(
sess, old_conv_op)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
new_sess = equalize_model(sess, conv_op.inputs[0].op.name, 'Relu_1')
new_conv_op = new_sess.graph.get_operation_by_name('conv2d/Conv2D')
self.assertFalse(BiasUtils.is_bias_none(new_conv_op))
conv_bias_data_after_fold = BiasUtils.get_bias_as_numpy_data(
new_sess, new_conv_op)
for i in range(len(conv_bias_data_before_fold)):
self.assertTrue(
conv_bias_data_before_fold[i] <= conv_bias_data_after_fold[i])
sess.close()
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 _call_mo_correct_bias(corrected_model: tf.compat.v1.Session,
layer_name: str,
bias_correction: libpymo.BiasCorrection,
bias_shape: int):
"""
helper to perform bias correction using cpp backend
:param corrected_model: active tensorflow session with corrected model as tf.compat.v1.Session
:param layer_name: name of the layer to be bias corrected
:param bias_correction: bias correction inputs
:param bias_shape: shape of bias associated with the layer
:return: None, updates bias for the given layer
"""
bias_tensor = libpymo.TensorParamBiasCorrection()
layer_to_be_corrected = corrected_model.graph.get_operation_by_name(
layer_name)
with corrected_model.graph.as_default():
assert (layer_to_be_corrected.type
in ['Conv2D', 'DepthwiseConv2dNative', 'MatMul'])
if BiasUtils.is_bias_none(layer_to_be_corrected):
bias_tensor.data = np.zeros(bias_shape)
else:
# read bias from given op
bias_tensor.data = BiasUtils.get_bias_as_numpy_data(
corrected_model, layer_to_be_corrected)
# perform bias correction
bias_correction.correctBias(bias_tensor)
# this api updates bias or adds bias add to layer if not present
BiasUtils.update_bias_for_op(corrected_model,
layer_to_be_corrected,
np.array(bias_tensor.data))
def bias_correction_single_layer_empirical(dataset: tf.data.Dataset):
""" perform bias correction on one layer """
# load a model
tf.keras.backend.clear_session()
_ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
sess = tf.compat.v1.keras.backend.get_session()
# input parameters for bias correction
# populate required parameters in two data types QuantParams and BiasCorrectParams
quant_params = QuantParams(quant_mode='tf_enhanced',
round_mode='nearest',
use_cuda=True,
ops_to_ignore=None)
bias_correction_params = BiasCorrectionParams(
batch_size=1,
num_quant_samples=10,
num_bias_correct_samples=10,
input_op_names=['input_1'],
output_op_names=['fc1000/Softmax'])
with sess.as_default():
# initialize model with zero bias
sess = BiasUtils.initialize_model_with_bias(
sess, bias_correction_params.input_op_names,
bias_correction_params.output_op_names)
# pick a layer for bias correction
example_conv_layer = sess.graph.get_operation_by_name(
'res2a_branch2a/Conv2D')
# invoke bias correction of one layer
BiasCorrection.bias_correction_per_layer(
reference_model=sess,
corrected_model=sess,
bias_correct_params=bias_correction_params,
layer_name_to_be_corrected=example_conv_layer.name,
quant_params=quant_params,
data_set=dataset)
sess.close()
def bias_correction_single_layer_analytical():
""" perform analytical bias correction on one layer """
# load a model
tf.keras.backend.clear_session()
_ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
sess = tf.compat.v1.keras.backend.get_session()
# input parameters for bias correction
# populate required parameters in two data types QuantParams and BiasCorrectParams
quant_params = QuantParams(quant_mode='tf_enhanced',
round_mode='nearest',
use_cuda=True,
ops_to_ignore=None)
with sess.as_default():
# initialize model with zero bias
sess = BiasUtils.initialize_model_with_bias(sess, ['input_1'],
['fc1000/Softmax'])
# pick a layer for bias correction
example_conv_layer = sess.graph.get_operation_by_name(
'res2a_branch2a/Conv2D')
# get candidate conv bns in the model
convs_bn_activation_info_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, ['input_1'], ['fc1000/Softmax'])
# make sure to pick example_conv_layer that has a bn op associated with it
if example_conv_layer in convs_bn_activation_info_dict.keys():
preceding_bn_layer_info = convs_bn_activation_info_dict[
example_conv_layer]
# invoke analytical bias correction on this layer
BiasCorrection.analytical_bias_correction_per_layer(
sess, example_conv_layer, preceding_bn_layer_info,
quant_params)
sess.close()
def test_initialize_with_bias_with_detached_ops(self):
""" Test that initialize with bias only affects valid ops """
tf.compat.v1.reset_default_graph()
sess = tf.compat.v1.Session()
inputs = tf.keras.Input(shape=(32, 32, 3,))
conv1 = tf.keras.layers.Conv2D(32, (3, 3), use_bias=False)(inputs)
_ = tf.keras.layers.Conv2D(16, (2, 2), activation=tf.nn.tanh, use_bias=False)(conv1)
_ = tf.keras.layers.Conv2D(8, (2, 2), activation=tf.nn.tanh)(conv1)
graph_editor.detach_inputs(sess.graph.get_operation_by_name('conv2d_1/Conv2D'))
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
# Check that outputs of conv2d and conv2d_1 have no biases
self.assertTrue(sess.graph.get_operation_by_name('conv2d/Conv2D').outputs[0].consumers()[0].type != 'BiasAdd')
self.assertTrue(sess.graph.get_operation_by_name('conv2d_1/Conv2D').outputs[0].consumers()[0].type != 'BiasAdd')
sess = BiasUtils.initialize_model_with_bias(sess, ['input_1'], ['conv2d_2/BiasAdd'])
# Check that conv2d has a bias inserted but not conv2d_1
self.assertTrue(sess.graph.get_operation_by_name('conv2d/Conv2D').outputs[0].consumers()[0].type == 'BiasAdd')
self.assertTrue(sess.graph.get_operation_by_name('conv2d_1/Conv2D').outputs[0].consumers()[0].type != 'BiasAdd')
def test_bias_correction_model_tf_with_no_bias(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), use_bias=False)(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), use_bias=False)(relu_2)
_ = 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')
w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
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])
# save and load the updated graph after high bias fold update
n_sess = save_and_load_graph('./test_update', sess)
conv_op = n_sess.graph.get_operation_by_name('conv2d/Conv2D')
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(
n_sess,
bias_correction_params,
quant_params,
dataset,
perform_only_empirical_bias_corr=False)
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()
def test_analytical_empirical_bias_correction(self):
"""
Test bn based bias correction hybrid with a user passed in dictionary of conv and bn after cle.
"""
# create a custom model
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
32,
32,
3,
), name="inputs")
conv_op = tf.keras.layers.Conv2D(
32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(inputs)
conv1_op = tf.keras.layers.Conv2D(
32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(conv_op)
bn_op = tf.keras.layers.BatchNormalization(
fused=True,
beta_initializer='random_uniform',
gamma_initializer='random_uniform',
moving_mean_initializer='random_uniform',
moving_variance_initializer='random_uniform')(conv1_op,
training=False)
conv2_op = tf.keras.layers.Conv2D(
32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(bn_op)
bn_op2 = tf.keras.layers.BatchNormalization(
fused=True,
beta_initializer='random_uniform',
gamma_initializer='random_uniform',
moving_mean_initializer='random_uniform',
moving_variance_initializer='random_uniform')(conv2_op,
training=False)
relu_1 = tf.nn.relu(bn_op2)
conv6_op = tf.keras.layers.Conv2D(32, (3, 3))(relu_1)
_ = tf.nn.relu(conv6_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
sess.run(init)
output_op = sess.graph.get_operation_by_name('Relu_1')
input_op_name = "inputs"
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)
input_tensor = sess.graph.get_tensor_by_name('inputs:0')
shape = input_tensor.shape
dataset = np.random.rand(1, shape[1], shape[2], shape[3])
# store conv bns info
conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, [input_op_name], [output_op.name])
# perform CLE
new_sess = equalize_model(sess, input_op_name, output_op.name)
conv_with_bn_op = new_sess.graph.get_operation_by_name(
'conv2d_1/Conv2D')
old_bias_as_numpy = BiasUtils.get_bias_as_numpy_data(
new_sess, conv_with_bn_op)
# perform bias correction and check analytical is performed.
with unittest.mock.patch(
'aimet_tensorflow.bias_correction.iter_first_x'
) as iter_first_x:
iter_first_x.return_value = [dataset]
with unittest.mock.patch(
'aimet_tensorflow.bias_correction.BiasCorrection.analytical_bias_correction_per_layer',
return_value=sess
) as mocked_analytical_bias_correction_per_layer:
updated_sess = BiasCorrection.correct_bias(
new_sess,
bias_corr_input,
quant_params,
dataset,
conv_bn_dict=conv_bn_dict,
perform_only_empirical_bias_corr=False)
self.assertEqual(
mocked_analytical_bias_correction_per_layer.call_count, 3)
sess.close()
new_sess.close()
def _fold_given_auto_selected_batch_norms(
sess: tf.compat.v1.Session,
layer_pairs: List[PairType]) -> tf.compat.v1.Session:
"""
Fold a given set of batch_norm layers into conv layers
:param sess: tf.compat.v1.Session
:param layer_pairs: pair of conv and bn layers
:return: new session with updated graph
"""
with sess.graph.as_default():
for pair in layer_pairs:
conv_linear, batchnorm, is_batch_norm_second = pair
assert conv_linear.type in [
'Conv2D', 'DepthwiseConv2dNative', 'MatMul'
]
# check flag
is_bias_valid = False
if not BiasUtils.is_bias_none(conv_linear):
is_bias_valid = True
bn_params = _get_bn_params(sess, batchnorm.op)
weight_tensor = _get_weight_tensor_transpose_reshape(
sess, conv_linear)
bias_tensor = _get_bias_tensor(sess, conv_linear)
bias = libpymo.fold(bn_params, weight_tensor, bias_tensor,
is_bias_valid, is_batch_norm_second)
# converting back to TF format [kh, kw, Nic, Noc] before updating weight tensor value
if conv_linear.type == 'DepthwiseConv2dNative':
# Depthwise conv layers in TF have outputs(Noc) set to 1.
# we send in format [Nic, Noc, kh, kw]
numpy_weight_reshaped = np.reshape(
weight_tensor.data, weight_tensor.shape).transpose(
(2, 3, 0, 1))
elif conv_linear.type == 'MatMul':
# o, i - convert to i , o
numpy_weight_reshaped = np.reshape(
weight_tensor.data,
[weight_tensor.shape[0], weight_tensor.shape[1]
]).transpose(1, 0)
else:
# conv2D case
# we sent in format [Noc, Nic, kh, kw]
numpy_weight_reshaped = np.reshape(
weight_tensor.data, weight_tensor.shape).transpose(
(2, 3, 1, 0))
WeightTensorUtils.update_tensor_for_op(sess, conv_linear,
numpy_weight_reshaped)
# remove bn op
BNUtils.skip_bn_op(sess, batchnorm.op, batchnorm.in_tensor,
batchnorm.out_tensor)
# update bias tensor, even in case there was no existing bias add op in given conv2D op.
bias_tensor_shape = [weight_tensor.shape[0]]
numpy_bias_reshaped = np.reshape(bias, bias_tensor_shape)
BiasUtils.update_bias_for_op(sess, conv_linear,
numpy_bias_reshaped)
# we edited the graph, so we should load and save for the metagraph associated with the session to be updated
after_bn_fold_sess = save_and_load_graph('./temp_bn_fold', sess)
return after_bn_fold_sess
def reduce_conv2d(sess: tf.compat.v1.Session,
op_tensor_tuple: Tuple[Op, List[tf.Tensor]], op_mask) -> (str, tf.Operation, tf.Operation):
"""
Conv2D module reducer
:param sess: current tf.compat.v1.Session
:param op_tensor_tuple: tuple containing the op to reduce, and a list of input tensors to the op
:param op_mask: Mask containing information on input and output channels to winnow
"""
tf_op = op_tensor_tuple[0].get_module()
padding = str(tf_op.get_attr("padding"), "utf-8")
strides = tf_op.get_attr("strides")
# Remove last part of conv op's name, or else we end up with names like Conv2D/Conv2D/Conv2D... if the same conv op
# is reduced multiple times
last_slash = op_tensor_tuple[0].dotted_name.rfind('/')
name = "reduced_" + op_tensor_tuple[0].dotted_name[:last_slash]
kernel_product = op_tensor_tuple[0].get_param_product('kernel')
kernel_size = kernel_product.shape.as_list()[:2]
# Depthwise Conv2d always has output dim of 1. Only slice out input channel dimensions.
if op_tensor_tuple[0].type == 'DepthwiseConv2dNative':
# Check to make sure input channel and output channel sizes are identical
# Format is expected to be NHWC, so channels is the last index in shape array
if tf_op.inputs[0].shape.as_list()[-1] != tf_op.outputs[0].shape.as_list()[-1]:
raise NotImplementedError('Reducing depthwise conv2d with differing input and output channel sizes not '
'supported')
output_dim = None
else:
output_dim = 3
reduced_weights = _get_reduced_params(sess=sess,
product=kernel_product,
mask=op_mask,
input_dim=2,
output_dim=output_dim)
reduced_weights_init = tf.constant_initializer(reduced_weights, verify_shape=True)
bias_product = op_tensor_tuple[0].get_param_product('bias')
reduced_bias = None
if bias_product:
use_bias = True
reduced_bias = _get_reduced_params(sess=sess,
product=bias_product,
mask=op_mask,
input_dim=None,
output_dim=0)
reduced_bias_init = tf.constant_initializer(reduced_bias, verify_shape=True)
else:
use_bias = False
reduced_bias_init = 'zeros'
output_ch_masks = op_mask.output_channel_masks
output_ch_indices_to_reduce = get_zero_positions_in_binary_mask(output_ch_masks[0])
filters = len(output_ch_masks[0]) - len(output_ch_indices_to_reduce)
# Check for regularization in the kernel
kernel_tensor = kernel_product.tensor_dict[kernel_product.consumers[0]]
kernel_regularizer = _get_kernel_regularizer(kernel_tensor)
if op_tensor_tuple[0].type == 'DepthwiseConv2dNative':
new_tensor = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size,
strides=strides[1:3],
padding=padding,
use_bias=use_bias,
bias_initializer=reduced_bias_init,
kernel_initializer=reduced_weights_init,
kernel_regularizer=kernel_regularizer,
name=name)(op_tensor_tuple[1][0])
else:
new_tensor = tf.keras.layers.Conv2D(filters=filters,
kernel_size=kernel_size,
strides=strides[1:3],
padding=padding,
use_bias=use_bias,
bias_initializer=reduced_bias_init,
kernel_initializer=reduced_weights_init,
kernel_regularizer=kernel_regularizer,
name=name)(op_tensor_tuple[1][0])
module = new_tensor.op
if use_bias:
module = module.inputs[0].op
WeightTensorUtils.update_tensor_for_op(sess, module, reduced_weights)
if use_bias:
BiasUtils.update_bias_for_op(sess, module, reduced_bias)
return name, new_tensor.op, module
def test_equalize_model_multi_input(self):
"""
Test bn fold with multiple input nodes
"""
tf.compat.v1.reset_default_graph()
tf.set_random_seed(0)
input1 = tf.keras.Input(name='input1', shape=(10, 10, 3))
input2 = tf.keras.Input(name='input2', shape=(12, 12, 3))
x1 = tf.keras.layers.Conv2D(
8, (1, 1),
name='conv1a',
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(input1)
x2 = tf.keras.layers.Conv2D(
8, (3, 3),
name='conv1b',
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(x1)
x3 = tf.keras.layers.Conv2D(
8, (3, 3),
name='conv1c',
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(input2)
x4 = tf.keras.layers.Conv2D(
8, (3, 3),
name='conv1d',
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(x3)
x = tf.keras.layers.add([x2, x4])
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(x)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv2_op)
_ = tf.nn.relu(bn_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
conv_1b_before_equalize = sess.graph.get_operation_by_name(
'conv1b/Conv2D')
conv_1b_bias_data_before_fold = BiasUtils.get_bias_as_numpy_data(
sess, conv_1b_before_equalize)
conv_1d_before_equalize = sess.graph.get_operation_by_name(
'conv1d/Conv2D')
conv_1d_bias_data_before_fold = BiasUtils.get_bias_as_numpy_data(
sess, conv_1d_before_equalize)
new_sess = equalize_model(sess, ["input1", "input2"], 'Relu')
conv_1b_after_equalize = new_sess.graph.get_operation_by_name(
'conv1b/Conv2D')
conv_1b_bias_data_after_fold = BiasUtils.get_bias_as_numpy_data(
new_sess, conv_1b_after_equalize)
conv_1d_after_equalize = new_sess.graph.get_operation_by_name(
'conv1d/Conv2D')
conv_1d_bias_data_after_fold = BiasUtils.get_bias_as_numpy_data(
new_sess, conv_1d_after_equalize)
for i in range(len(conv_1b_bias_data_after_fold)):
self.assertTrue(conv_1b_bias_data_after_fold[i] <=
conv_1b_bias_data_before_fold[i])
for i in range(len(conv_1d_bias_data_after_fold)):
self.assertTrue(conv_1d_bias_data_after_fold[i] <=
conv_1d_bias_data_before_fold[i])
sess.close()
def test_high_bias_fold_custom_model(self):
"""
Test high bias fold with a custom model
"""
tf.compat.v1.reset_default_graph()
tf.set_random_seed(0)
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))(relu_1)
bn_op_2 = tf.keras.layers.BatchNormalization(fused=True)(conv2_op)
conv3_op = tf.keras.layers.Conv2D(32, (3, 3))(bn_op_2)
relu_2 = tf.nn.relu(conv3_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
np.random.seed(0)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
start_op_names = conv_op.inputs[0].op.name
output_op_names = relu_2.op.name
# fold batchnorm layers
after_bn_fold_sess, folded_pairs = fold_all_batch_norms(
sess, start_op_names, output_op_names)
# replace any ReLU6 layers with ReLU
graph_util = GraphSearchUtils(after_bn_fold_sess.graph, start_op_names,
output_op_names)
after_relu_replace_sess = graph_util.find_and_replace_relu6_with_relu(
after_bn_fold_sess)
# perform cross-layer scaling on applicable layer sets
after_cls_sess, cls_set_info_list = CrossLayerScaling.scale_model(
after_relu_replace_sess, start_op_names, output_op_names)
# we want to validate that after high bias fold, bias for conv is >= bias before high bias fold.
conv_op = after_cls_sess.graph.get_operation_by_name('conv2d_2/Conv2D')
before_high_bias_fold_bias_data = BiasUtils.get_bias_as_numpy_data(
after_cls_sess, conv_op)
# perform high-bias fold
after_hbf_sess = HighBiasFold.bias_fold(after_cls_sess, folded_pairs,
cls_set_info_list)
# read updated bias value
conv_op = after_hbf_sess.graph.get_operation_by_name('conv2d_2/Conv2D')
high_bias_folded_bias_data = BiasUtils.get_bias_as_numpy_data(
after_hbf_sess, conv_op)
for i in range(len(before_high_bias_fold_bias_data)):
# folded bias should be greater than previous bias
self.assertTrue(high_bias_folded_bias_data[i] >=
before_high_bias_fold_bias_data[i])
sess.close()
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 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)
def analytical_bias_correction_per_layer(
corrected_model: tf.compat.v1.Session,
layer: tf.Operation,
preceeding_bn_layer_info: ConvBnInfoType,
quant_params: QuantParams,
is_first_conv: bool = False) -> tf.compat.v1.Session:
"""
Perform bn based bias correction (analytical bc).
:param corrected_model: active tensorflow session for corrected model
:param layer: conv/linear layer to be corrected
:param preceeding_bn_layer_info: corresponding preceeding bn/ activation info
:param quant_params: Quantization specific params from user
:param is_first_conv: flag to indicate if it's the first conv layer
:return: None, updates corrected_model in place
"""
layer = corrected_model.graph.get_operation_by_name(layer.name)
# get bn param and quantized weights from conv for this layer
bias_tensor, weight_tensor = BiasCorrection._get_conv_linear_params(
corrected_model, layer)
quantized_weight = BiasCorrection._get_quantized_weights(
weight_tensor, quant_params)
bn_params = libpymo.BnParamsBiasCorr()
activation_type = libpymo.ActivationType.noActivation
if preceeding_bn_layer_info:
input_tf_bn_op_name = preceeding_bn_layer_info.input_bn.get_module(
).name
bn_op = corrected_model.graph.get_operation_by_name(
input_tf_bn_op_name)
bn_params = BiasCorrection._get_bn_params(corrected_model, bn_op)
if preceeding_bn_layer_info.in_activation_type == ActivationType.relu:
activation_type = libpymo.ActivationType.relu
elif preceeding_bn_layer_info.in_activation_type == ActivationType.relu6:
activation_type = libpymo.ActivationType.relu6
elif preceeding_bn_layer_info.in_activation_type == ActivationType.no_activation:
activation_type = libpymo.ActivationType.noActivation
else:
assert (0, 'Unknown activation type',
preceeding_bn_layer_info.in_activation_type)
else:
if is_first_conv:
# for the first conv layer case, we use gamma = 1 and beta = 0
shape = weight_tensor.shape[1]
bn_params.gamma = np.ones(shape)
bn_params.beta = np.zeros(shape)
else:
assert 0, "layer info is None and is not first conv layer"
# need to invoke cpp api for bn based bias correction
biasCorrection = libpymo.BnBasedBiasCorrection()
biasCorrection.correctBias(bias_tensor, quantized_weight,
weight_tensor, bn_params, activation_type)
# this api updates bias or adds bias add to layer if not present
layer = corrected_model.graph.get_operation_by_name(layer.name)
BiasUtils.update_bias_for_op(corrected_model, layer,
np.array(bias_tensor.data))
logger.info('Completed analytical bias correction for layer %s',
layer.name)
def correct_bias(reference_model: tf.compat.v1.Session,
bias_correct_params: BiasCorrectionParams,
quant_params: QuantParams,
data_set: tf.data.Dataset,
conv_bn_dict: Union[Dict[tf.Operation, ConvBnInfoType],
None] = None,
perform_only_empirical_bias_corr: bool = True):
"""
Top level function for bias correction
:param reference_model: active tf.compat.v1.Session for the model to be corrected.
:param bias_correct_params: input params for bias correction
:param quant_params: QuantParams type with params for quantization simulation for bias correction.
:param data_set: input data set
:param conv_bn_dict: Dict of conv and bn with activation info. If None, the function looks for it.
This can be obtained on the model with bns and convs using
BiasCorrection.find_all_convs_bn_with_activation() api.
:param perform_only_empirical_bias_corr: a flag to indicate only empirical bias correction is to be performed.
:return: updated session with corrected bias for given ops
"""
# one time initialization of all layers with bias param
reference_model = BiasUtils.initialize_model_with_bias(
reference_model, bias_correct_params.input_op_names,
bias_correct_params.output_op_names)
# Create a copy of the model as reference model
corrected_model = save_and_load_graph('./temp_meta_path',
reference_model)
# get all ordered convs/ linears and skip gradient ops
ordered_conv_linears = get_ordered_conv_linears(
reference_model, bias_correct_params.input_op_names,
bias_correct_params.output_op_names)
# Get conv2D, depthwise with preceding BN ops info for analytical bias correction
# if user has not passed any dictionary
if conv_bn_dict is None:
convs_bn_activation_info_dict = BiasCorrection.find_all_convs_bn_with_activation(
reference_model, bias_correct_params.input_op_names,
bias_correct_params.output_op_names)
else:
convs_bn_activation_info_dict = BiasCorrection.refresh_op_ref(
reference_model, conv_bn_dict)
# Perform analytical bias correction for first conv layer
# we always perform empirical bias correction for linear layers
if ordered_conv_linears:
if not perform_only_empirical_bias_corr and ordered_conv_linears[
0].type not in ['MatMul']:
first_conv = ordered_conv_linears.pop(0)
BiasCorrection.analytical_bias_correction_per_layer(
corrected_model,
first_conv,
None,
quant_params,
is_first_conv=True)
# for each candidate layer in an ordered list of conv/lieanr ops
# find the corresponding bn and activation info
for layer in ordered_conv_linears:
# if this layer is in selected patterns of convs with preceding BN op and
# if empirical flag is false
# perform analytical Bias correction
if layer in convs_bn_activation_info_dict.keys(
) and not perform_only_empirical_bias_corr:
preceding_bn_layer_info = convs_bn_activation_info_dict[layer]
BiasCorrection.analytical_bias_correction_per_layer(
corrected_model, layer, preceding_bn_layer_info,
quant_params)
else:
# stand-alone convs/ linears or when perform_only_empirical_bias_corr is set to True
# perform empirical bias correction
BiasCorrection.bias_correction_per_layer(
reference_model, corrected_model, bias_correct_params,
layer.name, quant_params, data_set)
logger.info('Completed bias correction')
return corrected_model
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 test_depthwise_custom(self):
""" test depthwise conv2d layer withput bias """
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
10,
10,
3,
))
x = tf.keras.layers.Conv2D(10, (1, 1))(inputs)
with tf.compat.v1.variable_scope("standalone_depthwise"):
x = tf.compat.v1.nn.depthwise_conv2d_native(
x,
tf.compat.v1.get_variable(
initializer=tf.random.truncated_normal(shape=(3, 3, 10,
1)),
name="depthwise_kernel"), [1, 1, 1, 1], 'VALID')
_ = tf.nn.relu(x)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
sess.run(init)
op_list = sess.graph.get_operations()
depthwise_conv_op = sess.graph.get_operation_by_name(
'standalone_depthwise/DepthwiseConv2dNative')
input_op = sess.graph.get_operation_by_name('input_1')
output_op = sess.graph.get_operation_by_name('Relu')
input_op_names = ['input_1']
output_op_names = [output_op.name]
batch_size = 1
num_samples = 10
np.random.seed(0)
shape = input_op.outputs[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(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)
assert (BiasUtils.is_bias_none(depthwise_conv_op))
new_sess = BiasCorrection.correct_bias(sess, bias_correction_params,
quant_params, dataset)
updated_conv_op = new_sess.graph.get_operation_by_name(
'standalone_depthwise/DepthwiseConv2dNative')
assert (not BiasUtils.is_bias_none(updated_conv_op))
sess.close()
new_sess.close()