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 cross_layer_equalization_auto():
""" perform auto cross layer equalization """
# load a model
tf.keras.backend.clear_session()
_ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
sess = tf.compat.v1.keras.backend.get_session()
# get starting op name to invoke api for cle
input_op_name = 'input_1'
output_op_name = 'fc1000/Softmax'
# Equalize a model with Batchnorms
# Performs BatchNorm fold, replacing Relu6 with Relu, Cross layer scaling and High bias fold
# use the new session returned for further evaluations on TF graph
with sess.as_default():
new_session = equalize_model(sess, input_op_name, output_op_name)
sess.close()
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_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 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 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_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_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()