def test_with_slim_bn_op(self):
"""
Test with Tf Slim BN op
:return:
"""
tf.compat.v1.reset_default_graph()
sess = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph())
inp = tf.compat.v1.placeholder(tf.float32, [1, 32, 32, 3])
net = slim.conv2d(inp, 32, [3, 3])
_ = slim.batch_norm(net, decay=.7, epsilon=.65, is_training=True)
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
with sess.graph.as_default():
bn_op = sess.graph.get_operation_by_name('BatchNorm/FusedBatchNormV3')
moving_mean = BNUtils.get_moving_mean_as_numpy_data(sess, bn_op)
moving_var = BNUtils.get_moving_variance_as_numpy_data(sess, bn_op)
beta = BNUtils.get_beta_as_numpy_data(sess, bn_op)
gamma = BNUtils.get_gamma_as_numpy_data(sess, bn_op)
# check the values read are equal to init values
expected_beta = np.zeros_like(beta)
expected_gamma = np.ones_like(gamma)
expected_mean = np.zeros_like(moving_mean)
expected_variance = np.ones_like(moving_var)
self.assertTrue(np.allclose(expected_beta, beta))
self.assertTrue(np.allclose(expected_gamma, gamma))
self.assertTrue(np.allclose(expected_mean, moving_mean))
self.assertTrue(np.allclose(expected_variance, moving_var))
def _create_new_op(
self, op_tensor_tuple: Tuple[Op,
List[tf.Tensor]]) -> List[tf.Tensor]:
"""
Given a tuple of an operation to mirror and the parent tensor, create the new op. Return a list of output
tensors from the new op.
:param op_tensor_tuple: Tuple containing (op to winnow, list of input tensors to the op)
:return: List of output tensors of the winnowed op
"""
switcher = {
"Conv2D": module_reducers.reduce_conv2d,
"DepthwiseConv2dNative": module_reducers.reduce_conv2d,
"MaxPool": module_reducers.reduce_maxpool,
"BatchNorm": module_reducers.reduce_batchnorm,
"FusedBatchNormV3": module_reducers.reduce_batchnorm,
"Relu": module_reducers.reduce_relu,
"Relu6": module_reducers.reduce_relu,
"AvgPool": module_reducers.reduce_avgpool,
"Tanh": module_reducers.reduce_tanh,
"Add": module_reducers.reduce_add,
"AddN": module_reducers.reduce_add,
"AddV2": module_reducers.reduce_add,
"Identity": module_reducers.reduce_identity,
"Dropout": module_reducers.reduce_dropout,
"Pad": module_reducers.reduce_pad,
"PadV2": module_reducers.reduce_pad,
"MirrorPad": module_reducers.reduce_pad,
"Minimum": module_reducers.reduce_min_max,
"Maximum": module_reducers.reduce_min_max,
"Downsample": module_reducers.reduce_downsample,
"Upsample2D": module_reducers.reduce_upsample2d,
"LeakyRelu": module_reducers.reduce_leaky_relu
}
reducer = switcher.get(op_tensor_tuple[0].type,
module_reducers.reduce_default)
op_mask = self._op_to_mask_dict[op_tensor_tuple[0]]
name, output_op_node, module = reducer(self._sess, op_tensor_tuple,
op_mask)
self._reduced_op_info[op_tensor_tuple[0]] = ReducedInfo(
name, output_op_node, module)
self._reduced_modules[op_tensor_tuple[0].get_module().name] = (module,
op_mask)
output_tensors = output_op_node.outputs
if op_tensor_tuple[0].type in ['FusedBatchNormV3', 'Dropout']:
output_tensors = output_tensors[:
1] # only get first output tensor, but in list form
# Remove winnowed bn ops from UPDATE_OPS if present
if op_tensor_tuple[0].type == 'FusedBatchNormV3':
BNUtils.remove_bn_op_from_update_ops(
self._sess, op_tensor_tuple[0].get_module())
return output_tensors
def get_bn_params_aimet_api(sess, bn_op):
"""
Helper to get param values from BN layer using AIMET api(s)
:param bn_op: BN layer
:return: beta, gamma, mean and vairance values extracted from BN layer
"""
beta = BNUtils.get_beta_as_numpy_data(sess, bn_op)
gamma = BNUtils.get_gamma_as_numpy_data(sess, bn_op)
moving_mean = BNUtils.get_moving_mean_as_numpy_data(sess, bn_op)
moving_var = BNUtils.get_moving_variance_as_numpy_data(sess, bn_op)
return [beta, gamma, moving_mean, moving_var]
def _get_bn_params(model, bn_layer) -> libpymo.BnParamsBiasCorr():
"""
get bn params for bn based bias correction
:param model: tf.compat.v1.Session type
:param bn_layer: tf.Operation type
:return: bn params as libpymo.BnParamsBiasCorr() type
"""
bn_params = libpymo.BnParamsBiasCorr()
bn_params.beta = BNUtils.get_beta_as_numpy_data(model,
bn_layer).reshape(-1)
bn_params.gamma = BNUtils.get_gamma_as_numpy_data(model,
bn_layer).reshape(-1)
return bn_params
def get_bn_params_for_bias_fold(sess: tf.compat.v1.Session, bn_op: tf.Operation, scaling_parameter: np.ndarray):
"""
:param sess: active tf.compat.v1.Session
:param bn_op: tf Operation type fused batchnorm op.
:param scaling_parameter: scaling param as np.ndarray
:return: bn_params as BNParamsHighBiasFold type.
"""
bn_params = libpymo.BNParamsHighBiasFold()
# Scaling gamma and beta parameter of batch norm layer
gamma = BNUtils.get_gamma_as_numpy_data(sess, bn_op).reshape(-1)
bn_params.gamma = np.divide(gamma, scaling_parameter)
beta = BNUtils.get_beta_as_numpy_data(sess, bn_op).reshape(-1)
bn_params.beta = np.divide(beta, scaling_parameter)
return bn_params
def create_batchnorm_params(my_op: Op):
""" Create products for fusedbatchnorm """
tf_op = my_op.get_module()
beta_tensor = BNUtils.get_beta_read_var_op_tensor(tf_op)
create_and_connect_product('beta', beta_tensor.shape, my_op, beta_tensor)
gamma_tensor = BNUtils.get_gamma_read_var_op_tensor(tf_op)
create_and_connect_product('gamma', gamma_tensor.shape, my_op, gamma_tensor)
moving_mean_tensor = BNUtils.get_moving_mean_read_var_op_tensor(tf_op)
create_and_connect_product('moving_mean', moving_mean_tensor.shape, my_op,
moving_mean_tensor)
moving_variance_tensor = BNUtils.get_moving_variance_read_var_op_tensor(tf_op)
create_and_connect_product('moving_variance', moving_variance_tensor.shape, my_op,
moving_variance_tensor)
def _get_bn_params(sess: tf.compat.v1.Session,
bn: tf.Operation) -> libpymo.BNParams():
"""
helper to populate BN params from given BN op, required for fold
:param sess: tf.compat.v1.Session type
:param bn: BatchNorm or a FusedBatch Norm op
:return: bn_params
"""
# make sure you define the session and graph scope before loading vars from graph.
with sess.graph.as_default():
# create BNParams type and populate
bn_params = libpymo.BNParams()
bn_params.beta = BNUtils.get_beta_as_numpy_data(sess, bn).reshape(-1)
bn_params.gamma = BNUtils.get_gamma_as_numpy_data(sess, bn).reshape(-1)
bn_params.runningMean = BNUtils.get_moving_mean_as_numpy_data(
sess, bn).reshape(-1)
bn_params.runningVar = BNUtils.get_moving_variance_as_numpy_data(
sess, bn).reshape(-1)
epsilon = BNUtils.get_epsilon(bn)
var = BNUtils.get_moving_variance_as_numpy_data(sess, bn).reshape(-1)
var_with_epsilon = var + epsilon
sigma = np.sqrt(var_with_epsilon)
# sigma = tf.sqrt(BNUtils.get_moving_variance_as_numpy_data(sess, bn).reshape(-1) + epsilon)
bn_params.runningVar = sigma # sess.run(sigma).reshape(-1)
return bn_params
def test_training_batchnorm(self):
""" Test BNUtils get_training() with both fused and non fused batchnorms, with all three training modes """
tf.compat.v1.reset_default_graph()
# Model with fused batchnorms
_ = keras_model_functional()
fused_bn_training_true_op = tf.compat.v1.get_default_graph().get_operation_by_name('batch_normalization/FusedBatchNormV3')
self.assertTrue(BNUtils.get_training(fused_bn_training_true_op))
self.assertTrue(isinstance(BNUtils.get_training(fused_bn_training_true_op), bool))
fused_bn_training_tensor_op = tf.compat.v1.get_default_graph().get_operation_by_name('scope_1/batch_normalization_1/cond/'
'FusedBatchNormV3_1')
training_tensor = tf.compat.v1.get_default_graph().get_tensor_by_name('is_training:0')
self.assertEqual(BNUtils.get_training(fused_bn_training_tensor_op), training_tensor)
fused_bn_training_false_op = tf.compat.v1.get_default_graph().get_operation_by_name('scope_1/batch_normalization_2/'
'FusedBatchNormV3')
self.assertFalse(BNUtils.get_training(fused_bn_training_false_op))
tf.compat.v1.reset_default_graph()
# Model with non fused batchnorms
_ = keras_model_functional_with_non_fused_batchnorms()
bn_training_true_op = tf.compat.v1.get_default_graph().get_operation_by_name('batch_normalization/batchnorm/mul_1')
self.assertTrue(BNUtils.get_training(bn_training_true_op))
self.assertTrue(isinstance(BNUtils.get_training(bn_training_true_op), bool))
bn_training_tensor_op = tf.compat.v1.get_default_graph().get_operation_by_name('scope_1/batch_normalization_1/batchnorm/'
'mul_1')
training_tensor = tf.compat.v1.get_default_graph().get_tensor_by_name('is_training:0')
self.assertEqual(BNUtils.get_training(bn_training_tensor_op), training_tensor)
bn_training_false_op = tf.compat.v1.get_default_graph().get_operation_by_name('scope_1/batch_normalization_2/batchnorm/'
'mul_1')
self.assertFalse(BNUtils.get_training(bn_training_false_op))
tf.compat.v1.reset_default_graph()
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_batchnorm(sess: tf.compat.v1.Session,
op_tensor_tuple: Tuple[Op, List[tf.Tensor]], op_mask) -> (str, tf.Operation, tf.Operation):
"""
Fused and non fused batchnorm 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
"""
beta_product = op_tensor_tuple[0].get_param_product('beta')
if beta_product:
use_beta = True
reduced_beta_init = tf.constant_initializer(_get_reduced_params(sess=sess,
product=beta_product,
mask=op_mask,
input_dim=0,
output_dim=None),
verify_shape=True)
else:
use_beta = False
reduced_beta_init = 'zeros'
gamma_product = op_tensor_tuple[0].get_param_product('gamma')
if gamma_product:
use_gamma = True
reduced_gamma_init = tf.constant_initializer(_get_reduced_params(sess=sess,
product=gamma_product,
mask=op_mask,
input_dim=0,
output_dim=None),
verify_shape=True)
else:
use_gamma = False
reduced_gamma_init = 'ones'
moving_mean_product = op_tensor_tuple[0].get_param_product('moving_mean')
reduced_mov_mean_init = tf.constant_initializer(_get_reduced_params(sess=sess,
product=moving_mean_product,
mask=op_mask,
input_dim=0,
output_dim=None),
verify_shape=True)
moving_variance_product = op_tensor_tuple[0].get_param_product('moving_variance')
reduced_mov_variance_init = tf.constant_initializer(_get_reduced_params(sess=sess,
product=moving_variance_product,
mask=op_mask,
input_dim=0,
output_dim=None),
verify_shape=True)
name = "reduced_" + op_tensor_tuple[0].dotted_name
# Get training attribute
# This will either be True, False, or a string representing a training_placeholder the original BN was using
training = BNUtils.get_training(op_tensor_tuple[0].get_module())
assert training is not None
is_fused = op_tensor_tuple[0].type == 'FusedBatchNormV3'
epsilon = BNUtils.get_epsilon(op_tensor_tuple[0].get_module())
momentum = BNUtils.get_momentum(op_tensor_tuple[0].get_module())
if momentum is not None:
new_tensor = tf.keras.layers.BatchNormalization(center=use_beta,
scale=use_gamma,
epsilon=epsilon,
momentum=momentum,
beta_initializer=reduced_beta_init,
gamma_initializer=reduced_gamma_init,
moving_mean_initializer=reduced_mov_mean_init,
moving_variance_initializer=reduced_mov_variance_init,
fused=is_fused,
name=name)(op_tensor_tuple[1][0], training=training)
else:
new_tensor = tf.keras.layers.BatchNormalization(center=use_beta,
scale=use_gamma,
epsilon=epsilon,
beta_initializer=reduced_beta_init,
gamma_initializer=reduced_gamma_init,
moving_mean_initializer=reduced_mov_mean_init,
moving_variance_initializer=reduced_mov_variance_init,
fused=is_fused,
name=name)(op_tensor_tuple[1][0], training=training)
module = new_tensor.op.inputs[0].op
return name, new_tensor.op, module