def test_prune_model_2_layers(self): model = mnist_torch_model.Net() # Create a layer database orig_layer_db = LayerDatabase(model, input_shape=(1, 1, 28, 28)) # Copy the db comp_layer_db = copy.deepcopy(orig_layer_db) conv1 = comp_layer_db.find_layer_by_name('conv1') conv2 = comp_layer_db.find_layer_by_name('conv2') pruner = SpatialSvdPruner() layer_db = pruner.prune_model(orig_layer_db, [ LayerCompRatioPair(conv1, Decimal(0.5)), LayerCompRatioPair(conv2, Decimal(0.5)) ], CostMetric.mac, trainer=None) conv1_a = layer_db.find_layer_by_name('conv1.0') conv1_b = layer_db.find_layer_by_name('conv1.1') self.assertEqual((5, 1), conv1_a.module.kernel_size) self.assertEqual(1, conv1_a.module.in_channels) self.assertEqual(2, conv1_a.module.out_channels) self.assertEqual((1, 5), conv1_b.module.kernel_size) self.assertEqual(2, conv1_b.module.in_channels) self.assertEqual(32, conv1_b.module.out_channels) conv2_a = layer_db.find_layer_by_name('conv2.0') conv2_b = layer_db.find_layer_by_name('conv2.1') self.assertEqual((5, 1), conv2_a.module.kernel_size) self.assertEqual(32, conv2_a.module.in_channels) self.assertEqual(53, conv2_a.module.out_channels) self.assertEqual((1, 5), conv2_b.module.kernel_size) self.assertEqual(53, conv2_b.module.in_channels) self.assertEqual(64, conv2_b.module.out_channels) self.assertTrue(isinstance(layer_db.model.conv1, torch.nn.Sequential)) self.assertTrue(isinstance(layer_db.model.conv2, torch.nn.Sequential)) for layer in layer_db: print("Layer: " + layer.name) print(" Module: " + str(layer.module)) print(layer_db.model)
def testSpatialSvd(self): torch.manual_seed(1) model = mnist_torch_model.Net() rounding_algo = unittest.mock.MagicMock() rounding_algo.round.side_effect = [ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9 ] mock_eval = unittest.mock.MagicMock() mock_eval.side_effect = [ 100, 90, 80, 70, 60, 50, 40, 30, 20, 10, 90, 80, 70, 60, 50, 40, 30, 20, 10, 50 ] layer_db = LayerDatabase(model, input_shape=(1, 1, 28, 28)) pruner = SpatialSvdPruner() comp_ratio_select_algo = GreedyCompRatioSelectAlgo( layer_db, pruner, SpatialSvdCostCalculator(), mock_eval, 20, CostMetric.mac, Decimal(0.5), 10, True, None, rounding_algo, True, bokeh_session=None) layer_selector = ConvNoDepthwiseLayerSelector() spatial_svd_algo = CompressionAlgo( layer_db, comp_ratio_select_algo, pruner, mock_eval, layer_selector, modules_to_ignore=[], cost_calculator=SpatialSvdCostCalculator(), use_cuda=next(model.parameters()).is_cuda) compressed_layer_db, stats = spatial_svd_algo.compress_model( CostMetric.mac, trainer=None) self.assertTrue( isinstance(compressed_layer_db.model.conv1, torch.nn.Sequential)) self.assertTrue( isinstance(compressed_layer_db.model.conv2, torch.nn.Sequential)) self.assertTrue(stats.per_layer_stats[0].compression_ratio <= 0.5) self.assertEqual(0.3, stats.per_layer_stats[1].compression_ratio) print("Compressed model:") print(compressed_layer_db.model) print(stats)
def create_spatial_svd_algo(cls, model: torch.nn.Module, eval_callback: EvalFunction, eval_iterations, input_shape: Tuple, cost_metric: CostMetric, params: SpatialSvdParameters, bokeh_session: BokehServerSession) -> CompressionAlgo: """ Factory method to construct SpatialSvdCompressionAlgo :param model: Model to compress :param eval_callback: Evaluation callback for the model :param eval_iterations: Evaluation iterations :param input_shape: Shape of the input tensor for model :param cost_metric: Cost metric (mac or memory) :param params: Spatial SVD compression parameters :param bokeh_session: The Bokeh Session to display plots :return: An instance of SpatialSvdCompressionAlgo """ # pylint: disable=too-many-locals # Rationale: Factory functions unfortunately need to deal with a lot of parameters # Create a layer database layer_db = LayerDatabase(model, input_shape) use_cuda = next(model.parameters()).is_cuda # Create a pruner pruner = SpatialSvdPruner() cost_calculator = SpatialSvdCostCalculator() comp_ratio_rounding_algo = RankRounder(params.multiplicity, cost_calculator) # Create a comp-ratio selection algorithm if params.mode == SpatialSvdParameters.Mode.auto: greedy_params = params.mode_params.greedy_params comp_ratio_select_algo = GreedyCompRatioSelectAlgo(layer_db, pruner, cost_calculator, eval_callback, eval_iterations, cost_metric, greedy_params.target_comp_ratio, greedy_params.num_comp_ratio_candidates, greedy_params.use_monotonic_fit, greedy_params.saved_eval_scores_dict, comp_ratio_rounding_algo, use_cuda, bokeh_session=bokeh_session) layer_selector = ConvNoDepthwiseLayerSelector() modules_to_ignore = params.mode_params.modules_to_ignore else: # Convert (module,comp-ratio) pairs to (layer,comp-ratio) pairs layer_comp_ratio_pairs = cls._get_layer_pairs(layer_db, params.mode_params.list_of_module_comp_ratio_pairs) comp_ratio_select_algo = ManualCompRatioSelectAlgo(layer_db, layer_comp_ratio_pairs, comp_ratio_rounding_algo, cost_metric=cost_metric) layer_selector = ManualLayerSelector(layer_comp_ratio_pairs) modules_to_ignore = [] # Create the overall Spatial SVD compression algorithm spatial_svd_algo = CompressionAlgo(layer_db, comp_ratio_select_algo, pruner, eval_callback, layer_selector, modules_to_ignore, cost_calculator, use_cuda) return spatial_svd_algo
def test_prune_layer(self): model = mnist_torch_model.Net() # Create a layer database orig_layer_db = LayerDatabase(model, input_shape=(1, 1, 28, 28)) # Copy the db comp_layer_db = copy.deepcopy(orig_layer_db) conv1 = comp_layer_db.find_layer_by_name('conv1') spatial_svd_pruner = SpatialSvdPruner() spatial_svd_pruner._prune_layer(orig_layer_db, comp_layer_db, conv1, 0.5, CostMetric.mac) conv1_a = comp_layer_db.find_layer_by_name('conv1.0') conv1_b = comp_layer_db.find_layer_by_name('conv1.1') self.assertEqual((5, 1), conv1_a.module.kernel_size) self.assertEqual(1, conv1_a.module.in_channels) self.assertEqual(2, conv1_a.module.out_channels) self.assertEqual((1, 5), conv1_b.module.kernel_size) self.assertEqual(2, conv1_b.module.in_channels) self.assertEqual(32, conv1_b.module.out_channels) self.assertTrue( isinstance(comp_layer_db.model.conv1, torch.nn.Sequential)) for layer in comp_layer_db: print("Layer: " + layer.name) print(" Module: " + str(layer.module)) print(comp_layer_db.model) # check the output shapes of two newly created split layers # first split layer output conv1_a_output = comp_layer_db.model.conv1[0](torch.rand(1, 1, 28, 28)) # second split layer output conv1_b_output = comp_layer_db.model.conv1[1](conv1_a_output) self.assertEqual(conv1_a.output_shape, list(conv1_a_output.shape)) self.assertEqual(conv1_b.output_shape, list(conv1_b_output.shape))
def test_select_per_layer_comp_ratios_with_spatial_svd_pruner(self): pruner = SpatialSvdPruner() eval_func = unittest.mock.MagicMock() rounding_algo = unittest.mock.MagicMock() eval_func.side_effect = [ 10, 20, 30, 40, 50, 60, 70, 80, 90, 11, 21, 31, 35, 40, 45, 50, 55, 60 ] rounding_algo.round.side_effect = [ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9 ] model = mnist_torch_model.Net() layer_db = LayerDatabase(model, input_shape=(1, 1, 28, 28)) selected_layers = [ layer for layer in layer_db if isinstance(layer.module, nn.Conv2d) ] layer_db.mark_picked_layers(selected_layers) # Instantiate child greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db, pruner, SpatialSvdCostCalculator(), eval_func, 20, CostMetric.mac, Decimal(0.4), 10, True, None, rounding_algo, False, bokeh_session=None) layer_comp_ratio_list, stats = greedy_algo.select_per_layer_comp_ratios( ) original_cost = SpatialSvdCostCalculator.compute_model_cost(layer_db) for layer in layer_db: if layer not in selected_layers: layer_comp_ratio_list.append(LayerCompRatioPair(layer, None)) compressed_cost = SpatialSvdCostCalculator.calculate_compressed_cost( layer_db, layer_comp_ratio_list, CostMetric.mac) actual_compression_ratio = compressed_cost.mac / original_cost.mac self.assertTrue( math.isclose(Decimal(0.3), actual_compression_ratio, abs_tol=0.8)) print('\n') for pair in layer_comp_ratio_list: print(pair)
def test_eval_scores_with_spatial_svd_pruner(self): pruner = SpatialSvdPruner() eval_func = unittest.mock.MagicMock() eval_func.side_effect = [ 90, 80, 70, 60, 50, 40, 30, 20, 10, 91, 81, 71, 61, 51, 41, 31, 21, 11 ] model = mnist_torch_model.Net() # Create a layer database layer_db = LayerDatabase(model, input_shape=(1, 1, 28, 28)) layer1 = layer_db.find_layer_by_name('conv1') layer2 = layer_db.find_layer_by_name('conv2') layer_db.mark_picked_layers([layer1, layer2]) # Instantiate child greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo( layer_db, pruner, SpatialSvdCostCalculator(), eval_func, 20, CostMetric.mac, 0.5, 10, True, None, None, True, bokeh_session=None) dict = greedy_algo._compute_eval_scores_for_all_comp_ratio_candidates() print() print(dict) self.assertEqual(90, dict['conv1'][Decimal('0.1')]) self.assertEqual(51, dict['conv2'][Decimal('0.5')]) self.assertEqual(21, dict['conv2'][Decimal('0.8')])