def channel_pruning_manual_mode():
sess = tf.compat.v1.Session()
# Construct graph
with sess.graph.as_default():
_ = VGG16(weights=None, input_shape=(224, 224, 3))
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
# Create random dataset
batch_size = 1
input_data = np.random.rand(100, 224, 224, 3)
dataset = tf.data.Dataset.from_tensor_slices(input_data)
dataset = dataset.batch(batch_size=batch_size)
# Pick two convs to compress as examples
block1_conv2_op = sess.graph.get_operation_by_name('block1_conv2/Conv2D')
block2_conv2_op = sess.graph.get_operation_by_name('block2_conv2/Conv2D')
list_of_module_comp_ratio_pairs = [
ModuleCompRatioPair(block1_conv2_op, 0.5),
ModuleCompRatioPair(block2_conv2_op, 0.5)
]
manual_params = ChannelPruningParameters.ManualModeParams(
list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
params = ChannelPruningParameters(
input_op_names=['input_1'],
output_op_names=['predictions/Softmax'],
data_set=dataset,
batch_size=32,
num_reconstruction_samples=50,
allow_custom_downsample_ops=False,
mode=ChannelPruningParameters.Mode.manual,
params=manual_params,
multiplicity=8)
# Single call to compress the model
results = ModelCompressor.compress_model(
sess,
working_dir=None,
eval_callback=evaluate_model,
eval_iterations=10,
input_shape=(32, 224, 224, 3),
compress_scheme=CompressionScheme.channel_pruning,
cost_metric=CostMetric.mac,
parameters=params)
compressed_model, stats = results
print(compressed_model)
print(stats) # Stats object can be pretty-printed easily
def compress_session(sess, compressible_ops):
"""
Compressed TF session
:param sess: Tf session
:param compressible_ops: layers to compress
:return: compressed session
"""
layer_a = sess.graph.get_operation_by_name(compressible_ops[0])
list_of_module_comp_ratio_pairs = [ModuleCompRatioPair(layer_a, 0.5)]
manual_params = SpatialSvdParameters.ManualModeParams(
list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
params = SpatialSvdParameters(input_op_names=['input_1'], output_op_names=['act_softmax/Softmax'],
mode=SpatialSvdParameters.Mode.manual, params=manual_params)
scheme = CompressionScheme.spatial_svd
metric = CostMetric.mac
# pylint: disable=unused-argument
def evaluate(sess, iterations, use_cuda):
return 1
sess, _ = ModelCompressor.compress_model(sess=sess,
working_dir="./",
eval_callback=evaluate,
eval_iterations=None,
input_shape=(1, 3, 224, 224),
compress_scheme=scheme,
cost_metric=metric,
parameters=params)
return sess
def spatial_svd_manual_mode():
sess = tf.compat.v1.Session()
# Construct graph
with sess.graph.as_default():
_ = VGG16(weights=None, input_shape=(224, 224, 3))
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
# Pick two convs to compress as examples
conv2d = sess.graph.get_operation_by_name('block1_conv1/Conv2D')
conv2d_1 = sess.graph.get_operation_by_name('block1_conv2/Conv2D')
# Specify the necessary parameters
manual_params = SpatialSvdParameters.ManualModeParams([
ModuleCompRatioPair(module=conv2d, comp_ratio=0.5),
ModuleCompRatioPair(module=conv2d_1, comp_ratio=0.4)
])
params = SpatialSvdParameters(input_op_names=['input_1'],
output_op_names=['predictions/Softmax'],
mode=SpatialSvdParameters.Mode.manual,
params=manual_params)
input_shape = (1, 3, 224, 224)
# Single call to compress the model
compr_model_sess, stats = ModelCompressor.compress_model(
sess=sess,
working_dir=str('./'),
eval_callback=evaluate_model,
eval_iterations=10,
input_shape=input_shape,
compress_scheme=CompressionScheme.spatial_svd,
cost_metric=CostMetric.mac,
parameters=params,
trainer=None)
print(stats) # Stats object can be pretty-printed easily
def test_channel_pruning_manual_resnet50_keras(self):
"""
Manual mode test for Keras Resnet50
"""
AimetLogger.set_level_for_all_areas(logging.INFO)
graph = tf.Graph()
with graph.as_default():
_ = ResNet50(weights=None, input_shape=(224, 224, 3))
init = tf.compat.v1.global_variables_initializer()
# Grow GPU memory as needed at the cost of fragmentation.
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True # pylint: disable=no-member
sess = tf.compat.v1.Session(graph=graph, config=config)
# predicted value of the model
y_hat = sess.graph.get_tensor_by_name('probs/Softmax:0')
with sess.graph.as_default():
# place holder for the labels
y = tf.compat.v1.placeholder(tf.int64,
shape=[None, 1000],
name='labels')
# prediction Op
correct_prediction = tf.equal(tf.argmax(y_hat, axis=1),
tf.argmax(y, axis=1))
# accuracy Op
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32),
name='accuracy')
# initialize all the variables
sess.run(init)
batch_size = 32
input_data = np.random.rand(100, 224, 224, 3)
dataset = tf.data.Dataset.from_tensor_slices(input_data)
dataset = dataset.batch(batch_size=batch_size)
conv2_block1_1_conv = sess.graph.get_operation_by_name(
'conv2_block1_1_conv/Conv2D')
conv2_block1_2_conv = sess.graph.get_operation_by_name(
'conv2_block1_2_conv/Conv2D')
conv3_block1_1_conv = sess.graph.get_operation_by_name(
'conv3_block1_1_conv/Conv2D')
conv3_block1_2_conv = sess.graph.get_operation_by_name(
'conv3_block1_2_conv/Conv2D')
list_of_module_comp_ratio_pairs = [
ModuleCompRatioPair(conv2_block1_2_conv, 0.5),
ModuleCompRatioPair(conv3_block1_2_conv, 0.5)
]
input_op_names = ['input_1']
output_op_names = ['probs/Softmax']
manual_params = aimet_tensorflow.defs.ChannelPruningParameters.ManualModeParams(
list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
params = aimet_tensorflow.defs.ChannelPruningParameters(
input_op_names=input_op_names,
output_op_names=output_op_names,
data_set=dataset,
batch_size=32,
num_reconstruction_samples=50,
allow_custom_downsample_ops=True,
mode=aimet_tensorflow.defs.ChannelPruningParameters.Mode.manual,
params=manual_params,
multiplicity=8)
# channels_last data format
input_shape = (32, 224, 224, 3)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [70.123, 70.124]
compr_model_sess, stats = ModelCompressor.compress_model(
sess=sess,
working_dir=None,
eval_callback=mocked_eval,
eval_iterations=1,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.
channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params)
print(stats)
all_ops = compr_model_sess.graph.get_operations()
for op in all_ops:
if op.type == 'Conv2D':
print(op.name)
# conv2_block1_2_conv
reduced_conv2_block1_1_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_conv2_block1_1_conv/Conv2D')
reduced_conv2_block1_2_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_conv2_block1_2_conv/Conv2D')
# conv3_block1_2_conv
reduced_conv3_block1_1_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_conv3_block1_1_conv/Conv2D')
reduced_conv3_block1_2_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_conv3_block1_2_conv/Conv2D')
# conv2_block1_1_conv after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv2_block1_1_conv.inputs[1]).shape,
(1, 1, 64, 32))
self.assertEqual(
sess.run(conv2_block1_1_conv.inputs[1]).shape, (1, 1, 64, 64))
# conv2_block1_1_conv after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv2_block1_2_conv.inputs[1]).shape,
(3, 3, 32, 64))
self.assertEqual(
sess.run(conv2_block1_2_conv.inputs[1]).shape, (3, 3, 64, 64))
# conv3_block1_1_conv after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv3_block1_1_conv.inputs[1]).shape,
(1, 1, 256, 64))
self.assertEqual(
sess.run(conv3_block1_1_conv.inputs[1]).shape, (1, 1, 256, 128))
# conv3_block1_2_conv after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv3_block1_2_conv.inputs[1]).shape,
(3, 3, 64, 128))
self.assertEqual(
sess.run(conv3_block1_2_conv.inputs[1]).shape, (3, 3, 128, 128))
# Test forward pass through the model with random data and labels
new_input_data = np.random.rand(32, 224, 224, 3)
labels = np.random.randint(low=2, size=(32, 1000))
accuracy_tensor = compr_model_sess.graph.get_tensor_by_name(
'accuracy:0')
input_tensor = compr_model_sess.graph.get_tensor_by_name('input_1:0')
label_tensor = compr_model_sess.graph.get_tensor_by_name('labels:0')
compr_model_sess.run(accuracy_tensor,
feed_dict={
input_tensor: new_input_data,
label_tensor: labels
})
# close original session
sess.close()
# close compressed model session
compr_model_sess.close()
# delete temp directory
shutil.rmtree(str('./temp_meta/'))
def test_spatial_svd_compress_manual(self):
"""
End to end manual mode spatial SVD using Resnet50 Keras model
:return:
"""
np.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.INFO)
graph = tf.Graph()
with graph.as_default():
_ = ResNet50(weights=None, input_shape=(224, 224, 3))
init = tf.compat.v1.global_variables_initializer()
# Grow GPU memory as needed at the cost of fragmentation.
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True # pylint: disable=no-member
sess = tf.compat.v1.Session(graph=graph, config=config)
with sess.graph.as_default():
# predicted value of the model
y_hat = sess.graph.get_tensor_by_name('probs/Softmax:0')
# place holder for the labels
y = tf.compat.v1.placeholder(tf.int64,
shape=[None, 1000],
name='labels')
# prediction Op
correct_prediction = tf.equal(tf.argmax(y_hat, axis=1),
tf.argmax(y, axis=1))
# accuracy Op
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32),
name='accuracy')
sess.run(init)
writer = tf.compat.v1.summary.FileWriter('./', sess.graph)
# make sure the learning_phase flag is False (inference mode)
learning_phase = sess.graph.get_tensor_by_name(
'keras_learning_phase/input:0')
self.assertFalse(sess.run(learning_phase))
input_shape = (1, 224, 224, 3)
conv2_block1_2_conv = sess.graph.get_operation_by_name(
'conv2_block1_2_conv/Conv2D')
conv3_block1_2_conv = sess.graph.get_operation_by_name(
'conv3_block1_2_conv/Conv2D')
list_of_module_comp_ratio_pairs = [
ModuleCompRatioPair(conv2_block1_2_conv, 0.5),
ModuleCompRatioPair(conv3_block1_2_conv, 0.4)
]
manual_params = aimet_tensorflow.defs.SpatialSvdParameters.ManualModeParams(
list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
params = aimet_tensorflow.defs.SpatialSvdParameters(
input_op_names=['input_1'],
output_op_names=['probs/Softmax'],
mode=aimet_tensorflow.defs.SpatialSvdParameters.Mode.manual,
params=manual_params,
multiplicity=8)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [87, 64]
results = ModelCompressor.compress_model(
sess=sess,
working_dir=None,
eval_callback=mocked_eval,
eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params,
trainer=None)
compr_model_sess, stats = results
print(stats)
# split ops for res2a_branch2b_Conv2D op
conv2_block1_2_conv_a = compr_model_sess.graph.get_operation_by_name(
'conv2_block1_2_conv_a/Conv2D')
conv2_block1_2_conv_b = compr_model_sess.graph.get_operation_by_name(
'conv2_block1_2_conv_b/Conv2D')
# split ops for res3a_branch2b_Conv2D op
conv3_block1_2_conv_a = compr_model_sess.graph.get_operation_by_name(
'conv3_block1_2_conv_a/Conv2D')
conv3_block1_2_conv_b = compr_model_sess.graph.get_operation_by_name(
'conv3_block1_2_conv_b/Conv2D')
# res2a_branch2b_Conv2D
self.assertEqual(
compr_model_sess.run(conv2_block1_2_conv_a.inputs[1]).shape,
(3, 1, 64, 48))
self.assertEqual(
compr_model_sess.run(conv2_block1_2_conv_b.inputs[1]).shape,
(1, 3, 48, 64))
self.assertEqual(
compr_model_sess.run(conv3_block1_2_conv_a.inputs[1]).shape,
(3, 1, 128, 80))
self.assertEqual(
compr_model_sess.run(conv3_block1_2_conv_b.inputs[1]).shape,
(1, 3, 80, 128))
# forward pass to the model with random data and labels
input_data = np.random.rand(32, 224, 224, 3)
labels = np.random.randint(low=2, size=(32, 1000))
accuracy_tensor = compr_model_sess.graph.get_tensor_by_name(
'accuracy:0')
input_tensor = compr_model_sess.graph.get_tensor_by_name('input_1:0')
label_tensor = compr_model_sess.graph.get_tensor_by_name('labels:0')
# make sure the learning_phase flag is False (inference mode)
learning_phase = compr_model_sess.graph.get_tensor_by_name(
'keras_learning_phase/input:0')
self.assertFalse(compr_model_sess.run(learning_phase))
compr_model_sess.run(accuracy_tensor,
feed_dict={
input_tensor: input_data,
label_tensor: labels
})
# close original session
sess.close()
# close compressed model session
compr_model_sess.close()
# delete temp directory
shutil.rmtree(str('./temp_meta/'))
def test_svd_followed_by_channel_pruning(self):
""" Test that a model can be run through spatial svd and then channel pruning """
sess = tf.compat.v1.Session()
with sess.graph.as_default():
_ = model_with_three_convs()
init = tf.compat.v1.global_variables_initializer()
sess.run(init)
conv2d_1 = sess.graph.get_operation_by_name('conv2d_1/Conv2D')
list_of_module_comp_ratio_pairs = [ModuleCompRatioPair(conv2d_1, 0.5)]
manual_params = aimet_tensorflow.defs.SpatialSvdParameters.ManualModeParams(
list_of_module_comp_ratio_pairs=list_of_module_comp_ratio_pairs)
input_op_names = ['input_1']
output_op_names = ['three_convs/Softmax']
params = aimet_tensorflow.defs.SpatialSvdParameters(
input_op_names=input_op_names,
output_op_names=output_op_names,
mode=aimet_tensorflow.defs.SpatialSvdParameters.Mode.manual,
params=manual_params,
multiplicity=1)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [87, 87]
input_shape = (1, 8, 8, 3)
sess, _ = ModelCompressor.compress_model(
sess=sess,
working_dir=None,
eval_callback=mocked_eval,
eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params,
trainer=None)
# Check that svd added these ops
_ = sess.graph.get_operation_by_name('conv2d_1_a/Conv2D')
_ = sess.graph.get_operation_by_name('conv2d_1_b/Conv2D')
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.5),
num_comp_ratio_candidates=4,
use_monotonic_fit=True,
saved_eval_scores_dict=None)
conv_to_ignore = sess.graph.get_operation_by_name('conv2d/Conv2D')
auto_params = aimet_tensorflow.defs.ChannelPruningParameters.AutoModeParams(
greedy_select_params=greedy_params,
modules_to_ignore=[conv_to_ignore])
dataset = np.random.rand(1, 1, 8, 8, 3)
dataset = tf.convert_to_tensor(dataset)
dataset = tf.data.Dataset.from_tensor_slices(dataset)
params = aimet_tensorflow.defs.ChannelPruningParameters(
input_op_names=input_op_names,
output_op_names=output_op_names,
data_set=dataset,
batch_size=32,
num_reconstruction_samples=50,
allow_custom_downsample_ops=False,
mode=aimet_tensorflow.defs.ChannelPruningParameters.Mode.auto,
params=auto_params,
multiplicity=1)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [0, .5, 1, 0, .5, 1, 0, .5, 1, 0, 1]
sess, _ = ModelCompressor.compress_model(
sess=sess,
working_dir=None,
eval_callback=mocked_eval,
eval_iterations=1,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.
channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params)
# Check that these ops were added by cp
_ = sess.graph.get_operation_by_name(
'reduced_reduced_conv2d_1_a/Conv2D')
_ = sess.graph.get_operation_by_name(
'reduced_reduced_conv2d_1_b/Conv2D')
_ = sess.graph.get_operation_by_name('reduced_conv2d_2/Conv2D')