def bias_correction_after_cle(dataset: tf.data.Dataset):
"""
Perform bias correction on a given model (mix of empirical and analytical) after
cross layer equalization.
:param dataset: Data passed by user as tf.Dataset type.
:return: None
"""
# 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():
# store conv bns info before performing CLE
conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess,
start_op_names=['input_1'],
output_op_names=['fc1000/Softmax'])
# perform CLE
sess_after_cle = equalize_model(sess,
start_op_names=['input_1'],
output_op_names=['fc1000/Softmax'])
# run empirical and analytical bias correction on the model
_new_session = BiasCorrection.correct_bias(
sess_after_cle,
bias_correction_params,
quant_params,
dataset,
conv_bn_dict=conv_bn_dict,
perform_only_empirical_bias_corr=False)
sess.close()
def bias_correction_empirical(dataset: tf.data.Dataset):
"""
Perform bias correction on a given model
:param dataset: Data passed by user as tf.Dataset type.
:return: None
"""
# 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():
# run bias correction on the model
_new_session = BiasCorrection.correct_bias(sess,
bias_correction_params,
quant_params, dataset)
sess.close()
def test_bn_based_bias_correction_layer_selection_bn_conv(self):
"""
Test layer selection for bn based bias correction
patterns:
BN -> Conv
"""
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)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv_op)
relu = tf.nn.relu(bn_op)
conv1_op = tf.keras.layers.Conv2D(32, (3, 3))(relu)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(conv1_op)
_ = tf.nn.relu(conv2_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
start_op = "inputs"
bn_conv_linear_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, start_op, 'Relu_1')
assert (len(bn_conv_linear_dict) == 1)
sess.close()
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 test_correct_bias_on_mnist_with_analytical_bc(self):
""" Test correct bias using mnist model and data (analytical bias correction)"""
def modified_parse(serialized_example):
"""
Modified mnist dataset parser to provide only image data, without labels
:param serialized_example:
:return: Input image and labels
"""
dim = 28
features = tf.parse_single_example(serialized_example,
features={'label': tf.FixedLenFeature([], tf.int64),
'image_raw': tf.FixedLenFeature([], tf.string)})
# Mnist examples are flattened. Since we aren't performing an augmentations
# these can remain flattened.
image = tf.decode_raw(features['image_raw'], tf.uint8)
image.set_shape([dim*dim])
# Convert from bytes to floats 0 -> 1.
image = tf.cast(image, tf.float32) / 255
return image
tf.compat.v1.reset_default_graph()
batch_size = 2
num_samples = 10
dataset = tf.data.TFRecordDataset([os.path.join('data', 'mnist', 'validation.tfrecords')]).repeat(1)
dataset = dataset.map(modified_parse, num_parallel_calls=batch_size)
dataset = dataset.batch(batch_size=batch_size)
quant_params = QuantParams()
bias_correction_params = BiasCorrectionParams(batch_size=batch_size,
num_quant_samples=num_samples,
num_bias_correct_samples=num_samples,
input_op_names=['reshape_input'],
output_op_names=['dense_1/BiasAdd'])
sess = load_model_from_meta(meta_path=os.path.join('models', 'mnist_save.meta'))
BiasCorrection.correct_bias(sess, bias_correction_params, quant_params, dataset,
perform_only_empirical_bias_corr=False)
self.assertTrue(1) # Add some actual error check
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_bn_conv_layer_selection_bn_relu_conv(self):
"""
Test layer selection code
patterns:
BN -> Relu -> conv
BN -> (no activation) -> Conv
"""
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)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(conv_op,
training=False)
relu = tf.nn.relu(bn_op)
conv1_op = tf.keras.layers.Conv2D(32, (3, 3))(relu)
relu_2 = tf.nn.relu(conv1_op)
bn_op_2 = tf.keras.layers.BatchNormalization(fused=True)(relu_2)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(bn_op_2, training=False)
_ = tf.nn.relu(conv2_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
start_op = "inputs"
bn_conv_linear_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, start_op, 'Relu_2')
assert (2 == len(bn_conv_linear_dict))
conv1_op = sess.graph.get_operation_by_name('conv2d_1/Conv2D')
assert bn_conv_linear_dict[
conv1_op].in_activation_type == ActivationType.relu
assert bn_conv_linear_dict[conv1_op].input_bn is not None
conv2_op = sess.graph.get_operation_by_name('conv2d_2/Conv2D')
assert bn_conv_linear_dict[
conv2_op].in_activation_type == ActivationType.no_activation
assert bn_conv_linear_dict[conv2_op].input_bn is not None
sess.close()
def test_get_output_data(self):
"""
Test get_output_data method
"""
tf.compat.v1.reset_default_graph()
sess = tf.compat.v1.Session(graph=tf.Graph())
input_op_names = ['input_1']
output_op_name = 'scope_1/conv2d_2/Conv2D'
with sess.graph.as_default():
_ = keras_model_functional()
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
data = np.random.rand(1, 32, 32, 3)
output = BiasCorrection._get_output_data(sess, input_op_names,
output_op_name, data)
self.assertEqual(output.shape[3], 8)
sess.close()
def test_layer_selection_bn_depthwise_conv(self):
"""
Test layer selection with depthwise_layer on a custom model
patterns:
BN -> Depthwise conv
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
10,
10,
3,
))
x = tf.keras.layers.Conv2D(10, (1, 1))(inputs)
bn_op = tf.keras.layers.BatchNormalization(fused=True)(x,
training=False)
# 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(
bn_op,
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)
start_op = "input_1"
bn_conv_linear_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, start_op, 'Relu')
depthwise_op = sess.graph.get_operation_by_name(
'standalone_depthwise/DepthwiseConv2dNative')
assert (1 == len(bn_conv_linear_dict))
assert (depthwise_op in bn_conv_linear_dict.keys())
sess.close()
def run_evaluation(args):
# Build graph definition
with tf.Graph().as_default():
# Create iterator
tf_records = glob(args.dataset_dir + '/validation*')
preprocessing_fn = preprocessing_factory.get_preprocessing(
args.model_name, is_training=False)
parse_function = wrap_preprocessing(preprocessing_fn,
height=args.image_size,
width=args.image_size,
num_classes=(1001 -
args.labels_offset),
labels_offset=args.labels_offset)
dataset = tf.data.TFRecordDataset(tf_records).repeat(1)
dataset = dataset.map(parse_function, num_parallel_calls=1).apply(
tf.contrib.data.batch_and_drop_remainder(args.batch_size))
iterator = dataset.make_initializable_iterator()
images, labels = iterator.get_next()
network_fn = nets_factory.get_network_fn(
args.model_name,
num_classes=(1001 - args.labels_offset),
is_training=False)
with tf.device('/cpu:0'):
images = tf.placeholder_with_default(images,
shape=(None, args.image_size,
args.image_size, 3),
name='input')
labels = tf.placeholder_with_default(labels,
shape=(None, 1001 -
args.labels_offset),
name='labels')
logits, end_points = network_fn(images)
confidences = tf.nn.softmax(logits, axis=1, name='confidences')
categorical_preds = tf.argmax(confidences,
axis=1,
name='categorical_preds')
categorical_labels = tf.argmax(labels,
axis=1,
name='categorical_labels')
correct_predictions = tf.equal(categorical_labels, categorical_preds)
top1_acc = tf.reduce_mean(tf.cast(correct_predictions, tf.float32),
name='top1-acc')
top5_acc = tf.reduce_mean(tf.cast(
tf.nn.in_top_k(predictions=confidences,
targets=tf.cast(categorical_labels, tf.int32),
k=5), tf.float32),
name='top5-acc')
saver = tf.train.Saver()
sess = tf.Session()
# Load model from checkpoint
if not args.ckpt_bn_folded:
saver.restore(sess, args.checkpoint_path)
else:
sess.run(tf.global_variables_initializer())
# Fold all BatchNorms before QuantSim
sess, folded_pairs = fold_all_batch_norms(sess, ['IteratorGetNext'],
[logits.name[:-2]])
if args.ckpt_bn_folded:
with sess.graph.as_default():
saver = tf.train.Saver()
saver.restore(sess, args.checkpoint_path)
else:
# Do Cross Layer Equalization and Bias Correction if not loading from a batchnorm folded checkpoint
sess = equalize_model(sess, ['input'], [logits.op.name])
conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(
sess, ['input'], [logits.op.name])
quant_params = QuantParams(quant_mode=args.quant_scheme)
bias_correction_dataset = tf.data.TFRecordDataset(tf_records).repeat(1)
bias_correction_dataset = bias_correction_dataset.map(
lambda x: parse_function(x)[0], num_parallel_calls=1).apply(
tf.contrib.data.batch_and_drop_remainder(args.batch_size))
bias_correction_params = BiasCorrectionParams(
batch_size=args.batch_size,
num_quant_samples=10,
num_bias_correct_samples=512,
input_op_names=['input'],
output_op_names=[logits.op.name])
sess = BiasCorrection.correct_bias(
reference_model=sess,
bias_correct_params=bias_correction_params,
quant_params=quant_params,
data_set=bias_correction_dataset,
conv_bn_dict=conv_bn_dict,
perform_only_empirical_bias_corr=True)
# Define eval_func to use for compute encodings in QuantSim
def eval_func(session, iterations):
cnt = 0
avg_acc_top1 = 0
session.run('MakeIterator')
while cnt < iterations or iterations == -1:
try:
avg_acc_top1 += session.run('top1-acc:0')
cnt += 1
except:
return avg_acc_top1 / cnt
return avg_acc_top1 / cnt
# Select the right quant_scheme
if args.quant_scheme == 'range_learning_tf':
quant_scheme = aimet_common.defs.QuantScheme.training_range_learning_with_tf_init
elif args.quant_scheme == 'range_learning_tf_enhanced':
quant_scheme = aimet_common.defs.QuantScheme.training_range_learning_with_tf_enhanced_init
elif args.quant_scheme == 'tf':
quant_scheme = aimet_common.defs.QuantScheme.post_training_tf
elif args.quant_scheme == 'tf_enhanced':
quant_scheme = aimet_common.defs.QuantScheme.post_training_tf_enhanced
else:
raise ValueError("Got unrecognized quant_scheme: " + args.quant_scheme)
# Create QuantizationSimModel
sim = QuantizationSimModel(
session=sess,
starting_op_names=['IteratorGetNext'],
output_op_names=[logits.name[:-2]],
quant_scheme=quant_scheme,
rounding_mode=args.round_mode,
default_output_bw=args.default_output_bw,
default_param_bw=args.default_param_bw,
config_file=args.quantsim_config_file,
)
# Run compute_encodings
sim.compute_encodings(eval_func,
forward_pass_callback_args=args.encodings_iterations)
# Run final evaluation
sess = sim.session
top1_acc = eval_func(sess, -1)
print('Avg accuracy Top 1: {}'.format(top1_acc))
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_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()
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 for a single layer api
"""
# create a custom model
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
32,
32,
3,
), name="inputs")
conv1_op = tf.keras.layers.Conv2D(
32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(inputs)
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))(bn_op)
conv3_op = tf.keras.layers.Conv2D(32, (3, 3))(conv2_op)
conv4_op = tf.keras.layers.Conv2D(32, (3, 3))(conv3_op)
conv5_op = tf.keras.layers.Conv2D(32, (3, 3))(conv4_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')(conv5_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])
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:
with unittest.mock.patch(
'aimet_tensorflow.bias_correction.BiasCorrection.bias_correction_per_layer',
return_value=sess) as mocked_bias_correction_per_layer:
updated_sess = BiasCorrection.correct_bias(
sess,
bias_corr_input,
quant_params,
dataset,
perform_only_empirical_bias_corr=False)
assert mocked_bias_correction_per_layer.called
assert mocked_analytical_bias_correction_per_layer.called
self.assertEqual(
mocked_analytical_bias_correction_per_layer.call_count, 3)
self.assertEqual(mocked_bias_correction_per_layer.call_count,
3) # conv 3,4,5
sess.close()
updated_sess.close()
def test_bn_based_bias_correction_single_layer_functions_invoked(self):
"""
Test bn based bias correction for a single layer api methods invoked correctly
"""
# create a custom model
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
32,
32,
3,
), name="inputs")
conv1_op = tf.keras.layers.Conv2D(
32, (3, 3),
kernel_initializer=tf.random_uniform_initializer(-1, 1),
bias_initializer='random_uniform')(inputs)
conv2_op = tf.keras.layers.Conv2D(32, (3, 3))(conv1_op)
_ = tf.nn.relu(conv2_op)
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
sess.run(init)
conv_op = sess.graph.get_operation_by_name('conv2d/Conv2D')
output_op = sess.graph.get_operation_by_name('Relu')
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)
shape = conv_op.inputs[0].shape
dataset = np.random.rand(1, shape[1], shape[2], shape[3])
bias_tensor, weight = BiasCorrection._get_conv_linear_params(
sess, conv_op)
q_weight = BiasCorrection._get_quantized_weights(weight, quant_params)
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:
with unittest.mock.patch(
'aimet_tensorflow.bias_correction.BiasCorrection.bias_correction_per_layer',
return_value=sess) as mocked_bias_correction_per_layer:
updated_sess = BiasCorrection.correct_bias(
sess,
bias_corr_input,
quant_params,
dataset,
perform_only_empirical_bias_corr=False)
# check if api(s) are invoked
assert mocked_analytical_bias_correction_per_layer.called
called_args = mocked_analytical_bias_correction_per_layer.call_args
assert (called_args[1]['is_first_conv'] == True)
assert mocked_bias_correction_per_layer.called
sess.close()
updated_sess.close()
def main(args):
# configuration for efficient use of gpu
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
print('Loading srgan generator...')
gen_graph = tf.Graph()
with gen_graph.as_default():
gen_sess = tf.Session(config=config, graph=gen_graph)
with gen_sess.as_default():
srgan_generator = generator()
srgan_generator.load_weights(args.weights_path)
# sort files by filenames, assuming names match in both paths
lr_images_files = sorted(
glob.glob(os.path.join(args.images_path, '*LR.png')))
hr_images_files = sorted(
glob.glob(os.path.join(args.images_path, '*HR.png')))
# check if number of images align
if len(lr_images_files) != len(hr_images_files):
raise RuntimeError('length of image files doesn`t match,'
'need same number of images for both'
'low resolution and high resolution!')
image_files = (lr_images_files, hr_images_files)
# two list of metrics on all images
psnr_vals, ssim_vals = evaluate_session(gen_sess, image_files,
srgan_generator.input.name,
srgan_generator.output.name)
psnr_val = np.mean(psnr_vals)
ssim_val = np.mean(ssim_vals)
print(
f'Mean PSNR and SSIM for given images on original model are: [{psnr_val}, {ssim_val}]'
)
# TODO: use a better default dataset for compute encodings when not given by users
# use low resolution images if no representative lr data are provided
# use low and high resolution images if no representative lr and hr data are provided
if args.representative_datapath:
bc_lr_data = glob.glob(
os.path.join(args.representative_datapath, '*LR.png'))
comp_encodings_lr_data = glob.glob(
os.path.join(args.representative_datapath, '*LR.png'))
comp_encodings_hr_data = glob.glob(
os.path.join(args.representative_datapath, '*HR.png'))
else:
warnings.warn(
'No representative input data are given,'
'bias correction and computation of encodings will be done'
'on part of all of the low resolution images!')
bc_lr_data = lr_images_files
warnings.warn('No representative reference data are given,'
'computation of encodings will be done'
'on part of all of the high resolution images!')
comp_encodings_lr_data = lr_images_files
comp_encodings_hr_data = hr_images_files
comp_encodings_data = (comp_encodings_lr_data, comp_encodings_hr_data)
if args.cross_layer_equalization:
print('Applying cross layer equalization (CLE) to session...')
gen_sess = equalize_model(
gen_sess,
start_op_names=srgan_generator.input.op.name,
output_op_names=srgan_generator.output.op.name)
if args.bias_correction:
print('Applying Bias Correction (BC) to session...')
# the dataset being evaluated might have varying image sizes
# so right now only use batch size 1
batch_size = 1
num_imgs = len(bc_lr_data)
quant_params = QuantParams(use_cuda=args.use_cuda,
quant_mode=args.quant_scheme)
bias_correction_params = BiasCorrectionParams(
batch_size=batch_size,
num_quant_samples=min(num_imgs, args.num_quant_samples),
num_bias_correct_samples=min(num_imgs,
args.num_bias_correct_samples),
input_op_names=[srgan_generator.input.op.name],
output_op_names=[srgan_generator.output.op.name])
ds = make_dataset(bc_lr_data)
ds = ds.batch(batch_size)
gen_sess = BiasCorrection.correct_bias(gen_sess,
bias_correction_params,
quant_params, ds)
# creating quantsim object which inserts quantizer ops
sim = quantsim.QuantizationSimModel(
gen_sess,
starting_op_names=[srgan_generator.input.op.name],
output_op_names=[srgan_generator.output.op.name],
quant_scheme=args.quant_scheme,
default_output_bw=args.default_output_bw,
default_param_bw=args.default_param_bw)
# compute activation encodings
# usually achieves good results when data being used for computing
# encodings are representative of its task
partial_eval = partial(evaluate_session,
input_name=srgan_generator.input.name,
output_name='lambda_3/mul_quantized:0')
sim.compute_encodings(partial_eval, comp_encodings_data)
psnr_vals, ssim_vals = evaluate_session(sim.session,
image_files,
srgan_generator.input.name,
'lambda_3/mul_quantized:0',
output_dir=args.output_dir)
psnr_val = np.mean(psnr_vals)
ssim_val = np.mean(ssim_vals)
print(
f'Mean PSNR and SSIM for given images on quantized model are: [{psnr_val}, {ssim_val}]'
)
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 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)
_ = 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)
op_list = sess.graph.get_operations()
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]
op_list = sess.graph.get_operations()
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)
dense_conv_op = sess.graph.get_operation_by_name(
'single_residual/MatMul')
assert (BiasUtils.is_bias_none(dense_conv_op))
new_sess = BiasCorrection.correct_bias(sess, bias_correction_params,
quant_params, dataset)
updated_dense_conv_op = new_sess.graph.get_operation_by_name(
'single_residual/MatMul')
bias = BiasUtils.get_bias_as_numpy_data(new_sess,
updated_dense_conv_op)
assert (not BiasUtils.is_bias_none(updated_dense_conv_op))
sess.close()
new_sess.close()