def test_update_to_weight_tensor_with_load_var(self):
"""
tests update to weight tensor 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,))
_ = tf.keras.layers.Conv2D(32, (3, 3), kernel_initializer=tf.random_uniform_initializer(-1, 2))(inputs)
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_weights = WeightTensorUtils.get_tensor_as_numpy_data(sess, conv_op)
# add dummy weight tensor data
np.random.seed(0)
w_shape = WeightTensorUtils.get_tensor_shape(conv_op)
numpy_data = np.random.rand(3, w_shape[1], w_shape[2], w_shape[3])
# send in numpy data to overwrite previous value
WeightTensorUtils.update_tensor_for_op(sess, conv_op, numpy_data)
updated_weight_tensor = WeightTensorUtils.get_tensor_as_numpy_data(sess, conv_op)
# validate they are not the same
self.assertFalse(np.allclose(original_weights, updated_weight_tensor))
self.assertTrue(np.allclose(numpy_data, updated_weight_tensor))
sess.close()
def test_scale_cls_set_with_conv_layers_custom_model(self):
"""
Test scale_cls_set_with_conv_layers() on a custom model
"""
tf.set_random_seed(0)
_ = TestCrossLayerEqualization._custom_two_conv_layer_model()
init = tf.compat.v1.global_variables_initializer()
sess = tf.compat.v1.Session()
sess.run(init)
graph_util = GraphSearchUtils(tf.compat.v1.get_default_graph(),
"inputs", 'Relu')
_, layer_groups_as_tf_ops = graph_util.find_layer_groups_to_scale()
scaling_factors = CrossLayerScaling.scale_cls_set_with_conv_layers(
sess, layer_groups_as_tf_ops[0])
self.assertEqual(32, len(scaling_factors))
range_conv1_after_scaling = np.amax(np.abs(
WeightTensorUtils.get_tensor_as_numpy_data(
sess, layer_groups_as_tf_ops[0][0])),
axis=(2, 0, 1))
range_conv2_after_scaling = np.amax(np.abs(
WeightTensorUtils.get_tensor_as_numpy_data(
sess, layer_groups_as_tf_ops[0][1])),
axis=(3, 0, 1))
assert np.allclose(range_conv1_after_scaling,
range_conv2_after_scaling)
sess.close()
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 _select_inp_channels(layer: Layer, comp_ratio: float) -> list:
"""
:param layer: layer for which input channels to prune are selected.
:param comp_ratio: the ratio of costs after pruning has taken place
0 < comp_ratio <= 1.
:return: prune_indices: list of input channels indices to prune.
"""
assert layer.module.type == 'Conv2D'
weight_index = WeightTensorUtils.get_tensor_index_in_given_op(
layer.module)
weight_tensor = layer.model.run(layer.module.inputs[weight_index])
# Conv2d weight shape in TensorFlow [kh, kw, Nic, Noc]
# re order in the common shape [Noc, Nic, kh, kw]
weight_tensor = np.transpose(weight_tensor, (3, 2, 0, 1))
num_in_channels = weight_tensor.shape[1]
prune_indices = select_channels_to_prune(weight_tensor, comp_ratio,
num_in_channels)
return prune_indices
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 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 _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_update_weight_tensor_for_op(self):
""" Test update_weight_tensor_for_op() on custom conv op """
# get VGG16 model
tf.compat.v1.reset_default_graph()
tf.set_random_seed(0)
inputs = tf.keras.Input(shape=(
32,
32,
3,
), name="inputs")
conv_op = tf.keras.layers.Conv2D(32, (3, 3))(inputs)
_ = tf.nn.relu(conv_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')
initial_data = WeightTensorUtils.get_tensor_as_numpy_data(
sess, conv_op)
wt_data = initial_data + 2
# this is block1_conv1/Conv2D in VGG16
WeightTensorUtils.update_tensor_for_op(sess, conv_op, wt_data)
new_sess = aimet_tensorflow.utils.graph_saver.save_and_load_graph(
'./temp_conv_wt_updated', sess)
# check for if reroute was successful
# read op from conv op should be same as one defined by new variable type
conv_op = new_sess.graph.get_operation_by_name('conv2d/Conv2D')
new_wt_data = WeightTensorUtils.get_tensor_as_numpy_data(
new_sess, conv_op)
assert not np.allclose(initial_data, new_wt_data)
sess.close()
def get_weights(conv_module, sess):
"""
Returns the weights of a conv_module in a 2d matrix, where each column is an output channel.
:param sess: tf.compat.v1.Session
:param conv_module: convNd module
:return: 2d numpy array
"""
numpy_weight = WeightTensorUtils.get_tensor_as_numpy_data(
sess, conv_module)
numpy_weight = np.reshape(numpy_weight,
(numpy_weight.shape[3], numpy_weight.shape[2],
numpy_weight.shape[0], numpy_weight.shape[1]))
axis_0_length = numpy_weight.shape[0]
axis_1_length = np.prod(numpy_weight.shape[1:])
reshaped_weights = numpy_weight.reshape(int(axis_0_length),
int(axis_1_length))
return reshaped_weights
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_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_prune_layer(self):
""" Pruning single layer with 0.5 comp-ratio in MNIST"""
# 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
_ = mnist_tf_model.create_model(data_format='channels_last')
sess.run(tf.compat.v1.global_variables_initializer())
# Create a layer database
orig_layer_db = LayerDatabase(model=sess,
input_shape=(1, 28, 28, 1),
working_dir=None)
# Copy the db
comp_layer_db = copy.deepcopy(orig_layer_db)
conv1 = comp_layer_db.find_layer_by_name('conv2d/Conv2D')
# before the splitting
bias_op = get_succeeding_bias_op(conv1.module)
for consumer in bias_op.outputs[0].consumers():
self.assertEqual(consumer.name, "conv2d/Relu")
spatial_svd_pruner = SpatialSvdPruner()
spatial_svd_pruner._prune_layer(orig_layer_db, comp_layer_db, conv1,
0.5, CostMetric.mac)
conv2d_a_op = comp_layer_db.model.graph.get_operation_by_name(
'conv2d_a/Conv2D')
conv2d_b_op = comp_layer_db.model.graph.get_operation_by_name(
'conv2d_b/Conv2D')
conv2d_a_weight = WeightTensorUtils.get_tensor_as_numpy_data(
comp_layer_db.model, conv2d_a_op)
conv2d_b_weight = WeightTensorUtils.get_tensor_as_numpy_data(
comp_layer_db.model, conv2d_b_op)
conv1_a = comp_layer_db.find_layer_by_name('conv2d_a/Conv2D')
conv1_b = comp_layer_db.find_layer_by_name('conv2d_b/Conv2D')
# [Noc, Nic, kh, kw]
self.assertEqual([2, 1, 5, 1], conv1_a.weight_shape)
self.assertEqual([32, 2, 1, 5], conv1_b.weight_shape)
# after the splitting
bias_op = get_succeeding_bias_op(conv1_b.module)
for consumer in bias_op.outputs[0].consumers():
self.assertEqual(consumer.name, "conv2d/Relu")
# original layer should be not there in the database
self.assertRaises(
KeyError,
lambda: comp_layer_db.find_layer_by_name('conv2d/Conv2D'))
# check if the layer replacement is done correctly
orig_conv_op = comp_layer_db.model.graph.get_operation_by_name(
'conv2d/Conv2D')
bias_op = get_succeeding_bias_op(orig_conv_op)
# consumers list should be empty
consumers = [consumer for consumer in bias_op.outputs[0].consumers()]
self.assertEqual(len(consumers), 0)
# Check that weights loaded during svd pruning will stick after save and load
new_sess = save_and_load_graph('./temp_meta/', comp_layer_db.model)
conv2d_a_op = comp_layer_db.model.graph.get_operation_by_name(
'conv2d_a/Conv2D')
conv2d_b_op = comp_layer_db.model.graph.get_operation_by_name(
'conv2d_b/Conv2D')
conv2d_a_weight_after_save_load = WeightTensorUtils.get_tensor_as_numpy_data(
comp_layer_db.model, conv2d_a_op)
conv2d_b_weight_after_save_load = WeightTensorUtils.get_tensor_as_numpy_data(
comp_layer_db.model, conv2d_b_op)
self.assertTrue(
np.array_equal(conv2d_a_weight, conv2d_a_weight_after_save_load))
self.assertTrue(
np.array_equal(conv2d_b_weight, conv2d_b_weight_after_save_load))
tf.compat.v1.reset_default_graph()
sess.close()
new_sess.close()
# delete temp directory
shutil.rmtree(str('./temp_meta/'))
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 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 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 _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 test_bn_fold_with_linear_layer(self):
"""
test bn fold on matmul layer
Custom Model where BN layer is followed by MatMul layer
:return:
"""
tf.compat.v1.reset_default_graph()
inputs = tf.keras.Input(shape=(
1,
1,
4,
))
bn_op = tf.keras.layers.BatchNormalization(fused=True)(inputs,
training=False)
x = tf.keras.layers.Flatten()(bn_op)
_ = tf.keras.layers.Dense(2,
activation=tf.nn.relu,
name="linear_layer")(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()
linear_layer = sess.graph.get_operation_by_name('linear_layer/MatMul')
weight_before_fold = WeightTensorUtils.get_tensor_as_numpy_data(
sess, linear_layer)
input_op_name = 'input_1'
# get baseline output
np.random.seed(0)
input_tensor = sess.graph.get_tensor_by_name('input_1:0')
w_shape = input_tensor.shape
# tf 1.14 we do not have fused_batchnorm_1 in this case
bn_layer = sess.graph.get_operation_by_name(
'batch_normalization/FusedBatchNormV3')
numpy_data = np.random.rand(1, w_shape[1], w_shape[2], w_shape[3])
relu_op = sess.graph.get_operation_by_name('linear_layer/Relu')
baseline_output = sess.run(relu_op.outputs[0],
feed_dict={bn_layer.inputs[0]: numpy_data})
new_sess, pairs = fold_all_batch_norms(sess, input_op_name,
'linear_layer/Relu')
linear_layer = new_sess.graph.get_operation_by_name(
'linear_layer/MatMul')
weight_after_fold = WeightTensorUtils.get_tensor_as_numpy_data(
new_sess, linear_layer)
# check that weight got updated
self.assertFalse(
np.allclose(weight_before_fold, weight_after_fold, atol=1e-4))
# check outputs are close
linear_layer = new_sess.graph.get_operation_by_name(
'linear_layer/MatMul')
relu_op = new_sess.graph.get_operation_by_name('linear_layer/Relu')
# after bn removal, linear layer input is from flatten layer, that gets from input_1
after_fold_output = new_sess.run(
relu_op.outputs[0],
feed_dict={linear_layer.inputs[0].op.inputs[0]: numpy_data})
self.assertTrue(
np.allclose(baseline_output, after_fold_output, atol=1e-4))
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()