def test_calculate_weight_svd_cost_all_layers(self):
model = mnist_model.Net().to("cpu")
print(model)
layer_database = lad.LayerDatabase(model=model,
input_shape=(1, 1, 28, 28))
# Compress all layers by 50%
# Create a list of tuples of (layer, comp_ratio)
layer_ratio_list = []
for layer in layer_database:
if isinstance(layer.module, nn.Conv2d):
layer_ratio_list.append(
LayerCompRatioPair(layer, Decimal('0.5')))
else:
layer_ratio_list.append(
LayerCompRatioPair(layer, Decimal('0.5')))
compressed_cost = cc.WeightSvdCostCalculator.calculate_compressed_cost(
layer_database, layer_ratio_list, CostMetric.mac)
self.assertEqual(7031800, compressed_cost.mac)
def test_calculate_spatial_svd_cost_all_layers_given_ranks(self):
model = mnist_model.Net().to("cpu")
ld = lad.LayerDatabase(model=model, input_shape=(1, 1, 28, 28))
# Compress all layers by 50%
# Create a list of tuples of (layer, comp_ratio)
layer_rank_list = [(ld.find_layer_by_module(model.conv1), 2),
(ld.find_layer_by_module(model.conv2), 53),
(ld.find_layer_by_module(model.fc1), 385),
(ld.find_layer_by_module(model.fc2), 4)]
compressed_cost = cc.SpatialSvdCostCalculator.calculate_compressed_cost_given_ranks(
ld, layer_rank_list)
self.assertEqual(
5244960 + (3136 * 385 + 385 * 1024) + (1024 * 4 + 4 * 10),
compressed_cost.mac)
# Create a list of tuples of (layer, comp_ratio)
layer_rank_list = [(ld.find_layer_by_module(model.conv1), 2),
(ld.find_layer_by_module(model.conv2), 53),
(ld.find_layer_by_module(model.fc1), 385),
(ld.find_layer_by_module(model.fc2), None)]
compressed_cost = cc.SpatialSvdCostCalculator.calculate_compressed_cost_given_ranks(
ld, layer_rank_list)
self.assertEqual(5244960 + (3136 * 385 + 385 * 1024) + (1024 * 10),
compressed_cost.mac)
def test_calculate_spatial_svd_cost_all_layers(self):
model = mnist_model.Net().to("cpu")
print(model)
layer_database = lad.LayerDatabase(model=model,
input_shape=(1, 1, 28, 28))
model_cost = cc.SpatialSvdCostCalculator.compute_model_cost(
layer_database)
self.assertEqual(627200 + 10035200 + 3211264 + 10240, model_cost.mac)
# Compress all layers by 50%
# Create a list of tuples of (layer, comp_ratio)
layer_ratio_list = []
for layer in layer_database:
layer_ratio_list.append(LayerCompRatioPair(layer, Decimal(0.5)))
compressed_cost = cc.SpatialSvdCostCalculator.calculate_compressed_cost(
layer_database, layer_ratio_list, CostMetric.mac)
self.assertEqual(
5244960 + (3136 * 385 + 385 * 1024) + (1024 * 4 + 4 * 10),
compressed_cost.mac)
def test_total_model_cost(self):
logger.debug(self.id())
model = MnistSequentialModel().to("cpu")
layer_database = lad.LayerDatabase(model=model,
input_shape=(1, 1, 28, 28))
cost_calc = cc.CostCalculator()
network_cost = cost_calc.compute_model_cost(layer_database)
self.assertEqual(800 + 51200 + 3211264 + 10240, network_cost.memory)
self.assertEqual(627200 + 10035200 + 3211264 + 10240, network_cost.mac)
def test_calculate_channel_pruning_cost_all_layers(self):
model = mnist_model.Net().to("cpu")
print(model)
layer_database = lad.LayerDatabase(model=model,
input_shape=(1, 1, 28, 28))
# Compress all layers by 50%
# Create a list of tuples of (layer, comp_ratio)
layer_ratio_list = []
# Unfortunately in mnist we can only input channel prune conv2
for layer in layer_database:
if layer.module is model.conv2:
layer_ratio_list.append(
LayerCompRatioPair(layer, Decimal('0.5')))
else:
layer_ratio_list.append(LayerCompRatioPair(layer, None))
# Create the Input channel pruner
dataset_size = 1000
batch_size = 10
# create fake data loader with image size (1, 28, 28)
data_loader = self.create_fake_data_loader(dataset_size=dataset_size,
batch_size=batch_size)
pruner = InputChannelPruner(data_loader=data_loader,
input_shape=(1, 1, 28, 28),
num_reconstruction_samples=10,
allow_custom_downsample_ops=True)
cost_calculator = ChannelPruningCostCalculator(pruner)
compressed_cost = cost_calculator.calculate_compressed_cost(
layer_database, layer_ratio_list, CostMetric.mac)
self.assertEqual(8552704, compressed_cost.mac)
def test_calculate_spatial_svd_cost_linear_layer(self):
linear = nn.Linear(128, 256)
ld = lad.LayerDatabase(model=linear, input_shape=(1, 128))
layer = ld.find_layer_by_module(linear)
self.assertEqual(128,
cc.SpatialSvdCostCalculator.calculate_max_rank(layer))
comp_ratios_to_check = [1.0, 0.8, 0.75, 0.5, 0.25, 0.125]
original_cost = cc.CostCalculator.compute_layer_cost(layer)
self.assertEqual(128 * 256, original_cost.mac)
self.assertEqual(128 * 256, original_cost.memory)
for comp_ratio in comp_ratios_to_check:
rank = cc.SpatialSvdCostCalculator.calculate_rank_given_comp_ratio(
layer, comp_ratio, CostMetric.mac)
print('Rank = {}, for compression_ratio={}'.format(
rank, comp_ratio))
compressed_cost = cc.SpatialSvdCostCalculator.calculate_cost_given_rank(
layer, rank)
self.assertTrue(
math.isclose(compressed_cost.mac / original_cost.mac,
comp_ratio,
abs_tol=0.01))
# Higher level API
for comp_ratio in comp_ratios_to_check:
compressed_cost = cc.SpatialSvdCostCalculator.calculate_per_layer_compressed_cost(
layer, comp_ratio, CostMetric.mac)
self.assertTrue(
math.isclose(compressed_cost.mac / original_cost.mac,
comp_ratio,
abs_tol=0.01))
def __init__(self, model, run_model, run_model_iterations, input_shape,
compression_type, cost_metric, layer_selection_scheme,
**kw_layer_select_params):
"""Constructor for the Svd class
Constructs the Svd class from a set of options passed in at construction. The class takes
a number of named arguments which are detailed below.
:param model: The model which needs to be compressed
:param run_model: The evaluation function that needs to be passed for one forward pass
:param run_model_iterations: The number of iterations of forward pass for the run_model
:param input_shape: Shape of the inputs to the model
:param compression_type: Enum argument. Options available: svd , ssvd.
:param cost_metric: Enum argument. Options available: mac, memory
:param layer_selection_scheme: Enum argument. Options available: manual, top_n_layers, top_x_percent
:param kw_layer_select_params: Params for layer selection. Params depend on modes selected
1) If the layer_selection_scheme is manual then user has to specify the list of layers by using- layers_to_compress= [list of layers],
2) If the layer_selection_scheme is top_n_layers then the user has to specify the number of layers as num_layers= <number>
3) If the layer_selection_scheme is top_x_percent then the user has to specify percentage threshold by using percent_thresh= <number>
:return:
Todo: Add description
"""
# ------------------------------
# Initialize state
# ------------------------------
self._model = model
self._run_model = run_model
self._run_model_iterations = run_model_iterations
self._svd_lib_ref = pymo.GetSVDInstance()
self._network_cost = None
self._compression_type = compression_type
self._metric = cost_metric
# Layer selection related state
self._layer_selection_scheme = layer_selection_scheme
self._selected_layers = list()
self._use_cuda = self._is_model_on_gpu()
if self._use_cuda is True:
if not cuda.is_available():
raise ValueError(
'CUDA use was expected but CUDA is not available!')
# ------------------------------
# Creating the layer attribute database
self._layer_database = database.LayerDatabase(model=self._model,
input_shape=input_shape)
# picking layers for compression based on the scheme
ls.LayerSelectorDeprecated(layer_selection_scheme, cost_metric,
self._layer_database,
**kw_layer_select_params)
# Hack for now
if cost_metric == CostMetric.memory:
pymo_cost_metric = aimet_common.defs.CostMetric.memory
else:
pymo_cost_metric = aimet_common.defs.CostMetric.mac
pymo_utils.PymoSvdUtils.configure_layers_in_pymo_svd(
self._layer_database.get_selected_layers(), pymo_cost_metric,
self._svd_lib_ref)
logger.info("Selected layers: %s", self._selected_layers)
logger.info("Model is created on GPU : %s", self._use_cuda)