def test_compress_model_no_bias(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
model.conv2.bias = None
model.fc2.bias = None
c_model, stats = svd_intf.Svd.compress_model(model=model, run_model=mnist_model.evaluate,
run_model_iterations=1,
input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd,
cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.mac,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.manual,
rank_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.RankSelectionScheme.auto,
layers_to_compress=[model.conv2, model.fc2], num_rank_indices=20,
error_margin=100)
self.assertTrue(c_model.conv2[0].bias is None)
self.assertTrue(c_model.conv2[1].bias is None)
self.assertTrue(c_model.fc2[0].bias is None)
self.assertTrue(c_model.fc2[1].bias is None)
self.assertEqual(2, len(stats.per_rank_index[0].per_selected_layer))
self.assertEqual('conv2', stats.per_rank_index[0].per_selected_layer[0].layer_name)
self.assertEqual('fc2', stats.per_rank_index[0].per_selected_layer[1].layer_name)
def test_split_fc_layer_without_mo(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
with unittest.mock.patch('aimet_torch.layer_database.LayerDatabase'):
with unittest.mock.patch('aimet_torch.svd.layer_selector_deprecated.LayerSelectorDeprecated'):
svd = s.SvdImpl(model=model, run_model=None, run_model_iterations=1, input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd, cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.memory,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.top_n_layers, num_layers=2)
layer_attr = Layer(model.fc1, id(model.fc1), [3136, 1024, 1, 1])
svd._svd_lib_ref = create_autospec(pymo.Svd, instance=True)
split_weights = [np.zeros((400, model.fc1.in_features)).flatten().tolist(),
np.zeros((model.fc1.out_features, 400)).flatten().tolist()]
svd._svd_lib_ref.SplitLayerWeights.return_value = split_weights
split_biases = [np.zeros(400).flatten().tolist(),
np.zeros(model.fc1.out_features).flatten().tolist()]
svd._svd_lib_ref.SplitLayerBiases.return_value = split_biases
split_layer = svd_pruner_deprecated.DeprecatedSvdPruner
seq, layer_a_attr, layer_b_attr = split_layer.prune_layer(layer_attr, 400, svd_lib_ref=svd._svd_lib_ref)
self.assertEqual((400, model.fc1.in_features), seq[0].weight.shape)
self.assertEqual([400], list(seq[0].bias.shape))
self.assertEqual((model.fc1.out_features, 400), seq[1].weight.shape)
self.assertEqual([model.fc1.out_features], list(seq[1].bias.shape))
self.assertEqual(layer_a_attr.module, seq[0])
self.assertEqual(layer_b_attr.module, seq[1])
def test_spatial_svd_compress_auto_multi_input_model(self):
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = [(1, 1, 28, 28), (1, 1, 28, 28)]
model = ModelWithTwoInputs()
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=4)
auto_params = aimet_torch.defs.SpatialSvdParameters.AutoModeParams(greedy_params)
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.auto, auto_params)
mock_eval = FakeEvaluator(input_shape)
mock_eval.return_values = [0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25,
0.50, 0.50]
results = ModelCompressor.compress_model(model=model, eval_callback=mock_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,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertEqual(8, compressed_model.conv2[0].out_channels)
self.assertEqual((5, 1), compressed_model.conv2[0].kernel_size)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.25, abs_tol=0.01))
def test_weight_svd_compress_manual(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
manual_params = aimet_torch.defs.WeightSvdParameters.ManualModeParams(
[ModuleCompRatioPair(model.layer1[0].conv1, 0.5),
ModuleCompRatioPair(model.layer2[1].conv2, 0.4), ModuleCompRatioPair(model.fc, 0.4)])
params = aimet_torch.defs.WeightSvdParameters(aimet_torch.defs.WeightSvdParameters.Mode.manual, manual_params)
results = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.weight_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertTrue(isinstance(compressed_model.layer1[0].conv1, torch.nn.Sequential))
self.assertEqual(math.floor(64 * 64 * 3 * 3 * 0.5 / (64 + 64 * 3 * 3)),
compressed_model.layer1[0].conv1[0].out_channels)
self.assertTrue(isinstance(compressed_model.fc, torch.nn.Sequential))
self.assertEqual(math.floor(512 * 1000 * 0.4 / (512 + 1000)), compressed_model.fc[0].out_features)
def test_single_pytorch_module_mapping_to_many_onnx_nodes(self):
""" test onxx based utility to find mapping between onnx node names and io tensors
when more than one onnx node maps to the same torch module
"""
AimetLogger.set_level_for_all_areas(logging.DEBUG)
class TwoLayerLstmModel(torch.nn.Module):
"""
Model using torch.nn.LSTM module
"""
def __init__(self):
super(TwoLayerLstmModel, self).__init__()
self.lstm = torch.nn.LSTM(input_size=3, hidden_size=5, num_layers=3)
def forward(self, x, hx=None):
return self.lstm(x, hx)
model_name = 'multilayer_lstm'
model = TwoLayerLstmModel()
dummy_input = torch.randn(10, 1, 3)
torch.onnx.export(model, dummy_input, './data/' + model_name + '.onnx')
onnx_utils.OnnxSaver.set_node_names('./data/' + model_name + '.onnx', model, dummy_input)
onnx_model = onnx.load('./data/' + model_name + '.onnx')
lstm_nodes = [node for node in onnx_model.graph.node if node.op_type == 'LSTM']
self.assertEqual(3, len(lstm_nodes))
node_to_io_dict, _ = onnx_utils.OnnxSaver.get_onnx_node_to_io_tensor_names_map(onnx_model)
self.assertEqual(1, len(node_to_io_dict))
self.assertTrue(isinstance(node_to_io_dict['lstm'], list))
self.assertEqual(3, len(node_to_io_dict['lstm']))
def test_channel_pruning_manual(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
manual_params = ChannelPruningParameters.ManualModeParams([ModuleCompRatioPair(model.layer1[0].conv2, 0.3),
ModuleCompRatioPair(model.layer2[1].conv1, 0.5)])
params = ChannelPruningParameters(data_loader.train_loader, 5000,
True,
aimet_torch.defs.ChannelPruningParameters.Mode.manual,
manual_params, multiplicity=8)
compressed_model, stats = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params, visualization_url=None)
baseline_model_accuracy = stats.baseline_model_accuracy
compressed_best_model_accuracy = stats.compressed_model_accuracy
self.assertTrue(baseline_model_accuracy >= compressed_best_model_accuracy)
self.assertEqual(24, compressed_model.layer1[0].conv2.in_channels)
def test_spatial_svd_with_fine_tuning(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
# load trained MNIST model
data_loader = DataLoaderMnist(cuda=True, seed=1, shuffle=False, train_batch_size=64, test_batch_size=100)
input_shape = (1, 1, 28, 28)
model = torch.load(os.path.join('./', 'data', 'mnist_trained_on_GPU.pth'))
modules_to_ignore = [model.conv1]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.8),
num_comp_ratio_candidates=10,
use_monotonic_fit=True)
auto_params = aimet_torch.defs.SpatialSvdParameters.AutoModeParams(greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.auto, auto_params,
multiplicity=1)
results = ModelCompressor.compress_model(model, mnist_model.evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
trainer=Trainer(), visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.7, abs_tol=0.1))
def test_create_compressed_model(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
with unittest.mock.patch('aimet_torch.svd.layer_database.LayerDatabase'):
with unittest.mock.patch('aimet_torch.svd.layer_selector_deprecated.LayerSelectorDeprecated'):
svd = s.SvdImpl(model=model, run_model=None, run_model_iterations=1, input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd, cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.memory,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.top_n_layers, num_layers=2)
ls.LayerSelectorDeprecated._pick_compression_layers = create_autospec(ls.LayerSelectorDeprecated._pick_compression_layers)
layer_attr1 = Layer(model.fc2, id(model.fc2), model.fc2.weight.shape)
layer_attr1.parent_module = model
layer_attr1.var_name_of_module_in_parent = "fc2"
layer_attr1.output_shape = [0, 0, 1, 1]
layer_attr1.name = 'fc2'
layer_attr2 = Layer(model.conv2, id(model.conv2), model.conv2.weight.shape)
layer_attr2.parent_module = model
layer_attr2.var_name_of_module_in_parent = "conv2"
layer_attr2.name = 'conv2'
layer_attr1.output_shape = [0, 0, 14, 14]
ls.LayerSelectorDeprecated._pick_compression_layers.return_value = [layer_attr1, layer_attr2]
svd._compressible_layers = {id(model.conv2): layer_attr2,
id(model.fc2): layer_attr1}
ls.LayerSelectorDeprecated._perform_layer_selection(model)
svd._select_candidate_ranks(20)
svd_rank_pair_dict = {'conv2': (31,0), 'fc2': (9,0)}
c_model, c_layer_attr, _ = svd._create_compressed_model(svd_rank_pair_dict)
self.assertTrue(c_model is not model)
self.assertTrue(c_model.conv1 is not model.conv1)
self.assertTrue(c_model.conv2 is not model.conv2)
self.assertFalse(isinstance(svd._model, nn.Sequential))
self.assertEqual((9, 1024), c_model.fc2[0].weight.shape)
self.assertEqual([9], list(c_model.fc2[0].bias.shape))
self.assertEqual((10, 9), c_model.fc2[1].weight.shape)
self.assertEqual([10], list(c_model.fc2[1].bias.shape))
self.assertEqual((31, 32, 1, 1), c_model.conv2[0].weight.shape)
self.assertEqual([31], list(c_model.conv2[0].bias.shape))
self.assertEqual((64, 31, 5, 5), c_model.conv2[1].weight.shape)
self.assertEqual([64], list(c_model.conv2[1].bias.shape))
self.assertEqual(svd._model.conv1.weight.shape, c_model.conv1.weight.shape)
self.assertEqual(svd._model.fc1.weight.shape, c_model.fc1.weight.shape)
# Expect double the number of layers in layer_attr_list
self.assertEqual(4, len(c_layer_attr))
def test_channel_pruning_compress_auto_resnet(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
data_loader = ImageNetDataLoader(image_dir, image_size, batch_size, num_workers)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=False).to(torch.device('cuda'))
model.eval()
modules_to_ignore = [model.conv1,
model.layer2[0].downsample[0],
model.layer3[0].downsample[0],
model.layer4[0].downsample[0],
model.layer4[1].conv1,
model.layer4[1].conv2
]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=10,
use_monotonic_fit=True,
saved_eval_scores_dict=
'./data/resnet18_eval_scores.pkl')
auto_params = ChannelPruningParameters.AutoModeParams(greedy_params,
modules_to_ignore=modules_to_ignore)
# selecting single batch for reconstruction
# num_reconstruction_samples = 50
# 50 / 10 (samples_per_image) = 5 = batch size
params = ChannelPruningParameters(data_loader=data_loader.train_loader,
num_reconstruction_samples=50,
allow_custom_downsample_ops=True,
mode=aimet_torch.defs.ChannelPruningParameters.Mode.auto,
params=auto_params, multiplicity=8)
results = ModelCompressor.compress_model(model=model, eval_callback=evaluate, eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.channel_pruning,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertNotEqual(model, compressed_model)
self.assertTrue(0.6 < float(stats.mac_compression_ratio) < 0.65)
def test_spatial_svd_compress_auto_with_high_multiplicity(self):
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18()
modules_to_ignore = [model.conv1,
model.layer1[0].conv1, model.layer1[0].conv2,
model.layer1[1].conv1, model.layer1[1].conv2,
model.layer2[0].downsample[0],
model.layer3[0].conv1, model.layer3[0].conv2,
model.layer3[1].conv1, model.layer3[1].conv2,
model.layer3[0].downsample[0],
model.layer4[0].conv1, model.layer4[0].conv2,
model.layer4[0].downsample[0],
model.layer4[1].conv1,
model.layer4[1].conv2
]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.65),
num_comp_ratio_candidates=4)
auto_params = aimet_torch.defs.SpatialSvdParameters.AutoModeParams(greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.auto, auto_params,
multiplicity=64)
mock_eval = FakeEvaluator(input_shape)
mock_eval.return_values = [0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25, 0.75, 0.50, 0.25,
0.50, 0.50]
results = ModelCompressor.compress_model(model=model, eval_callback=mock_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,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertEqual(64, compressed_model.layer2[0].conv2[0].out_channels)
self.assertEqual((3, 1), compressed_model.layer2[0].conv2[0].kernel_size)
self.assertTrue(math.isclose(float(stats.mac_compression_ratio), 0.87, abs_tol=0.01))
def test_svd_manual_rank_sel_weight_svd_deprecated(self):
torch.cuda.empty_cache()
AimetLogger.set_level_for_all_areas(logging.DEBUG)
# load trained MNIST model
model = torch.load(os.path.join('./', 'data', 'mnist_trained_on_CPU.pth'))
compressed_model, stats = svd_intf.Svd.compress_model(model=model, run_model=mnist_model.evaluate,
run_model_iterations=1, input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd,
cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.mac,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.manual,
rank_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.RankSelectionScheme.manual,
layer_rank_list=[[model.conv2, 27]])
baseline_model_accuracy = stats.baseline_model_accuracy
compressed_best_model_accuracy = stats.compressed_model_accuracy
self.assertTrue(baseline_model_accuracy >= compressed_best_model_accuracy)
def test_add_pytorch_node_names_to_onnx_resnet(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
model_name = 'resnet18'
model = models.resnet18(pretrained=False)
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, './data/' + model_name + '.onnx')
onnx_utils.OnnxSaver.set_node_names('./data/' + model_name + '.onnx', model, dummy_input)
onnx_model = onnx.load('./data/' + model_name + '.onnx')
for node in onnx_model.graph.node:
if node.op_type in ('Conv', 'Gemm', 'MaxPool'):
self.assertTrue(node.name)
for in_tensor in node.input:
if in_tensor.endswith('weight'):
print("Checking " + in_tensor)
self.assertEqual(node.name, in_tensor[:-7])
def test_add_pytorch_node_names_to_onnx_ooo(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
model_name = 'out_of_order'
model = OutOfOrderModel()
dummy_input = torch.randn(1, 16, 20, 20)
torch.onnx.export(model, dummy_input, './data/' + model_name + '.onnx')
onnx_utils.OnnxSaver.set_node_names('./data/' + model_name + '.onnx', model, dummy_input)
onnx_model = onnx.load('./data/' + model_name + '.onnx')
for node in onnx_model.graph.node:
if node.op_type in ('Conv', 'Gemm', 'MaxPool'):
self.assertTrue(node.name)
for in_tensor in node.input:
if in_tensor.endswith('weight'):
print("Checking " + in_tensor)
self.assertEqual(node.name, in_tensor[:-7])
def test_weight_svd_compress_auto_high_multiplicity(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
torch.backends.cudnn.deterministic = True
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18(pretrained=True).to(torch.device('cuda'))
modules_to_ignore = [model.conv1,
model.layer2[0].downsample[0],
model.layer3[0].downsample[0],
model.layer4[0].downsample[0],
model.layer4[1].conv1,
model.layer4[1].conv2,
model.fc
]
greedy_params = aimet_common.defs.GreedySelectionParameters(target_comp_ratio=Decimal(0.7),
num_comp_ratio_candidates=4,
saved_eval_scores_dict=
'./data/resnet18_eval_scores.pkl')
rank_select = RankSelectScheme.greedy
auto_params = aimet_torch.defs.WeightSvdParameters.AutoModeParams(rank_select_scheme=rank_select,
select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_torch.defs.WeightSvdParameters(aimet_torch.defs.WeightSvdParameters.Mode.auto, auto_params,
multiplicity=64)
results = ModelCompressor.compress_model(model, evaluate, 10, input_shape,
aimet_common.defs.CompressionScheme.weight_svd,
cost_metric=aimet_common.defs.CostMetric.mac, parameters=params,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
self.assertFalse(isinstance(compressed_model.conv1, torch.nn.Sequential))
self.assertFalse(isinstance(compressed_model.fc, torch.nn.Sequential))
def test_setting_log_level_for_all_areas(self):
logger.info("*** test_setting_log_level_for_all_areas() ***\n")
svd_logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Svd)
quant_logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Quant)
util_logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Utils)
test_logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Test)
# The default logging level for all Log Areas defined in default_logging_config.json is used.
logger.info(
"Log at the default log level for all Log Areas defined in default_logging_config.json"
)
svd_logger.debug("Testing Debug")
svd_logger.info("Testing Info")
quant_logger.warning("Testing Warning")
quant_logger.error("Testing Error")
util_logger.critical("Testing Critical")
util_logger.info("Testing Info")
test_logger.critical("Testing Critical")
test_logger.critical("****************************************\n")
# Change the default log level for all areas
# Only CRITICAL level logs will be logged.
logger.info("Change the logging level for all Log Areas to WARNING")
AimetLogger.set_level_for_all_areas(logging.WARNING)
svd_logger.debug("Testing Debug")
svd_logger.info("Testing Info")
quant_logger.warning("Testing Warning")
quant_logger.error("Testing Error")
util_logger.critical("Testing Critical")
util_logger.info("Testing Info")
test_logger.critical("Testing Critical")
test_logger.critical("****************************************\n")
def test_spatial_svd_compress_manual(self):
torch.cuda.empty_cache()
torch.manual_seed(1)
numpy.random.seed(1)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
input_shape = (1, 3, 224, 224)
model = models.resnet18()
manual_params = aimet_torch.defs.SpatialSvdParameters.ManualModeParams(
[ModuleCompRatioPair(model.layer1[0].conv1, 0.5),
ModuleCompRatioPair(model.layer2[1].conv2, 0.4)])
params = aimet_torch.defs.SpatialSvdParameters(aimet_torch.defs.SpatialSvdParameters.Mode.manual, manual_params,
multiplicity=8)
# Only used in this test for baseline and final accuracy - essentially a don't care
mock_eval = MagicMock()
mock_eval.return_value = 50
results = ModelCompressor.compress_model(model=model, eval_callback=mock_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,
visualization_url=None)
compressed_model, stats = results
print(compressed_model)
print(stats)
# Check that indeed weight svd was applied to some layer
self.assertTrue(isinstance(compressed_model.layer1[0].conv1, torch.nn.Sequential))
self.assertEqual(48, compressed_model.layer1[0].conv1[0].out_channels)
self.assertEqual((3, 1), compressed_model.layer1[0].conv1[0].kernel_size)
self.assertEqual((1, 0), compressed_model.layer1[0].conv1[0].padding)
self.assertTrue(isinstance(compressed_model.layer2[1].conv2, torch.nn.Sequential))
self.assertEqual(80, compressed_model.layer2[1].conv2[0].out_channels)
def test_compress_model_with_stride(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
# Change the model to add a stride to conv2, and adjust the input dimensions of the next layer accordingly
model.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=(2, 2), stride=(2, 2))
model.fc1 = nn.Linear(3*3*64, 1024)
c_model, stats = svd_intf.Svd.compress_model(model=model, run_model=mnist_model.evaluate,
run_model_iterations=1,
input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd,
cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.mac,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.manual,
rank_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.RankSelectionScheme.auto,
layers_to_compress=[model.conv2, model.fc2], num_rank_indices=20,
error_margin=100)
self.assertEqual(2, len(stats.per_rank_index[0].per_selected_layer))
self.assertEqual('conv2', stats.per_rank_index[0].per_selected_layer[0].layer_name)
self.assertEqual('fc2', stats.per_rank_index[0].per_selected_layer[1].layer_name)
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_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_channel_pruning_auto_mobilenetv1(self):
"""
Auto mode test fot MobileNetv1 model
:return:
"""
AimetLogger.set_level_for_all_areas(logging.INFO)
graph = tf.Graph()
with graph.as_default():
_ = MobileNet(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('reshape_2/Reshape: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
_ = tf.reduce_mean(tf.cast(correct_prediction, tf.float32),
name='accuracy')
sess.run(init)
modules_to_ignore = []
all_ops = sess.graph.get_operations()
for op in all_ops:
if op.type == 'Conv2D':
modules_to_ignore.append(op)
del modules_to_ignore[2:4]
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.92),
num_comp_ratio_candidates=4,
use_monotonic_fit=True,
saved_eval_scores_dict=None)
auto_params = aimet_tensorflow.defs.ChannelPruningParameters.AutoModeParams(
greedy_select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
input_op_names = ['input_1']
output_op_names = ['reshape_2/Reshape']
dataset = np.random.rand(1, 1, 224, 224, 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=True,
mode=aimet_tensorflow.defs.ChannelPruningParameters.Mode.auto,
params=auto_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, 70.125, 70.123, 70.124, 70.125, 70.123, 70.124,
70.125
]
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)
conv_pw_1_op = sess.graph.get_operation_by_name('conv_pw_1/Conv2D')
conv_pw_2_op = sess.graph.get_operation_by_name('conv_pw_2/Conv2D')
conv_pw_3_op = sess.graph.get_operation_by_name('conv_pw_3/Conv2D')
reduced_conv_pw_1_op = compr_model_sess.graph.get_operation_by_name(
'reduced_conv_pw_1/Conv2D')
reduced_conv_pw_2_op = compr_model_sess.graph.get_operation_by_name(
'reduced_reduced_conv_pw_2/Conv2D')
reduced_conv_pw_3_op = compr_model_sess.graph.get_operation_by_name(
'reduced_conv_pw_3/Conv2D')
# conv_pw_1 after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv_pw_1_op.inputs[1]).shape,
(1, 1, 32, 32))
self.assertEqual(
sess.run(conv_pw_1_op.inputs[1]).shape, (1, 1, 32, 64))
# conv_pw_2 after and before compression weight shapes
# TODO: uncomment below check after dangling winnowed nodes in ConnectedGraph are removed
self.assertEqual(
compr_model_sess.run(reduced_conv_pw_2_op.inputs[1]).shape,
(1, 1, 32, 64))
self.assertEqual(
sess.run(conv_pw_2_op.inputs[1]).shape, (1, 1, 64, 128))
# conv_pw_3 after and before compression weight shapes
self.assertEqual(
compr_model_sess.run(reduced_conv_pw_3_op.inputs[1]).shape,
(1, 1, 64, 128))
self.assertEqual(
sess.run(conv_pw_3_op.inputs[1]).shape, (1, 1, 128, 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')
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_spatial_svd_compress_auto_with_finetuning(self):
"""
End to end test with MNIST model following fine tuning
:return:
"""
tf.compat.v1.set_random_seed(10)
AimetLogger.set_level_for_all_areas(logging.DEBUG)
# load the meta file
meta_path = os.path.join('models', 'mnist_save.meta')
sess = aimet_tensorflow.utils.graph_saver.load_model_from_meta(
meta_path)
# ignore first Conv2D op
conv2d = sess.graph.get_operation_by_name('conv1/Conv2D')
modules_to_ignore = [conv2d]
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.5),
num_comp_ratio_candidates=10,
use_monotonic_fit=True,
saved_eval_scores_dict=None)
auto_params = aimet_tensorflow.defs.SpatialSvdParameters.AutoModeParams(
greedy_select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_tensorflow.defs.SpatialSvdParameters(
input_op_names=['reshape_input'],
output_op_names=['dense_1/BiasAdd'],
mode=aimet_tensorflow.defs.SpatialSvdParameters.Mode.auto,
params=auto_params,
multiplicity=8)
input_shape = (1, 1, 28, 28)
compr_model_sess, stats = ModelCompressor.compress_model(
sess=sess,
working_dir=None,
eval_callback=evaluate,
eval_iterations=5,
input_shape=input_shape,
compress_scheme=aimet_common.defs.CompressionScheme.spatial_svd,
cost_metric=aimet_common.defs.CostMetric.mac,
parameters=params)
print(stats)
self.assertEqual(evaluate(compr_model_sess, 1, True),
float(stats.compressed_model_accuracy))
all_ops = compr_model_sess.graph.get_operations()
conv_ops = [op for op in all_ops if op.type == 'Conv2D']
self.assertEqual(len(conv_ops), 4)
self.assertTrue(
math.isclose(float(stats.mac_compression_ratio), 0.5, abs_tol=0.1))
# get the weights after fine tuning
conv2d_1_a_op = compr_model_sess.graph.get_operation_by_name(
'conv2_a/Conv2D')
conv2d_1_a_op_weights_before = conv2d_1_a_op.inputs[1].eval(
session=compr_model_sess)
# fine tune the model
# get the input and validation place holders
x = compr_model_sess.graph.get_tensor_by_name('reshape_input:0')
y = compr_model_sess.graph.get_tensor_by_name('labels:0')
cross_entropy = compr_model_sess.graph.get_tensor_by_name('xent:0')
with compr_model_sess.graph.as_default():
# new optimizer and back propagation Op
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=1e-3,
name='Adam_new')
train_step = optimizer.minimize(loss=cross_entropy,
name='train_step_new')
# initialize only uninitialized variables
# only needed when fine tuning, because we are adding new optimizer
graph_eval.initialize_uninitialized_vars(compr_model_sess)
mnist = input_data.read_data_sets(os.path.join(str('./'), 'data'),
one_hot=True)
for i in range(1):
batch = mnist.train.next_batch(batch_size=32, shuffle=True)
_, loss_val = compr_model_sess.run([train_step, cross_entropy],
feed_dict={
x: batch[0],
y: batch[1]
})
# get the weights after fine tuning
conv2d_1_a_op = compr_model_sess.graph.get_operation_by_name(
'conv2_a/Conv2D')
conv2d_1_a_op_weights_after = conv2d_1_a_op.inputs[1].eval(
session=compr_model_sess)
# weight should be different after one iteration
self.assertFalse(
np.allclose(conv2d_1_a_op_weights_before,
conv2d_1_a_op_weights_after))
# close original session
sess.close()
# close compressed model session
compr_model_sess.close()
# delete temp directory
shutil.rmtree(str('./temp_meta/'))
def test_channel_pruning_auto_resnet50(self):
"""
Auto mode test for Keras ResNet-50.
"""
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)
# create dataset using random numpy array
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)
# compressing only two Conv2D layers
modules_to_ignore = list()
all_ops = sess.graph.get_operations()
for op in all_ops:
if op.type == 'Conv2D':
modules_to_ignore.append(op)
del modules_to_ignore[5:7]
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.9),
num_comp_ratio_candidates=4,
use_monotonic_fit=False)
auto_params = aimet_tensorflow.defs.ChannelPruningParameters.AutoModeParams(
greedy_select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
input_op_names = ['input_1']
output_op_names = ['probs/Softmax']
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.auto,
params=auto_params,
multiplicity=1)
# channels_last data format
input_shape = (32, 224, 224, 3)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [70.123, 70.124, 70.125, 87, 64, 99, 22, 33]
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)
# updated conv2_block2_1_conv
updated_conv2_block2_1_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_reduced_conv2_block2_1_conv/Conv2D')
# updated conv2_block2_2_conv
updated_conv2_block2_2_conv = compr_model_sess.graph.get_operation_by_name(
'reduced_conv2_block2_2_conv/Conv2D')
# conv2_block2_1_conv after compression weight shapes
self.assertEqual(
compr_model_sess.run(updated_conv2_block2_1_conv.inputs[1]).shape,
(1, 1, 64, 16))
# conv2_block2_2_conv after compression weight shapes
self.assertEqual(
compr_model_sess.run(updated_conv2_block2_2_conv.inputs[1]).shape,
(3, 3, 16, 64))
# 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_auto(self):
"""
End to end auto 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)
all_ops = sess.graph.get_operations()
for op in all_ops:
print(op.name)
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)
# 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)
# compressing only two layers
modules_to_ignore = list()
all_ops = sess.graph.get_operations()
for op in all_ops:
if op.type == 'Conv2D':
modules_to_ignore.append(op)
del modules_to_ignore[5:7]
greedy_params = aimet_common.defs.GreedySelectionParameters(
target_comp_ratio=Decimal(0.7),
num_comp_ratio_candidates=3,
use_monotonic_fit=True,
saved_eval_scores_dict=None)
auto_params = aimet_tensorflow.defs.SpatialSvdParameters.AutoModeParams(
greedy_select_params=greedy_params,
modules_to_ignore=modules_to_ignore)
params = aimet_tensorflow.defs.SpatialSvdParameters(
input_op_names=['input_1'],
output_op_names=['probs/Softmax'],
mode=aimet_tensorflow.defs.SpatialSvdParameters.Mode.auto,
params=auto_params,
multiplicity=8)
mocked_eval = unittest.mock.MagicMock()
mocked_eval.side_effect = [50, 80, 30, 76, 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)
# 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/'))
import torch
import geffnet
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from tqdm import tqdm
import argparse
from aimet_torch import utils
from aimet_torch import cross_layer_equalization
from aimet_torch import batch_norm_fold
from aimet_common.defs import QuantScheme
from aimet_torch.quantsim import QuantizationSimModel
from aimet_torch.onnx_utils import onnx_pytorch_conn_graph_type_pairs
from aimet_common.utils import AimetLogger
import logging
AimetLogger.set_level_for_all_areas(logging.DEBUG)
onnx_pytorch_conn_graph_type_pairs.append([["Clip"], ["hardtanh"]])
def work_init(work_id):
seed = torch.initial_seed() % 2**32
random.seed(seed + work_id)
np.random.seed(seed + work_id)
def model_eval(data_loader, image_size, batch_size=64, quant=False):
def func_wrapper_quant(model, arguments):
top1_acc = 0.0
total_num = 0
idx = 0
iterations, use_cuda = arguments[0], arguments[1]
