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)