def prune_layer(original_layer, rank, svd_lib_ref): """ Splits a layer based on the splitting scheme :param original_layer: original layers attributes :param rank: rank pair for a given layer :param svd_lib_ref: Reference to Model optimization library :return: """ # Delegate to the right method to split the layer if isinstance(original_layer.module, nn.Conv2d): module_a, module_b = WeightSvdModuleSplitter.split_conv_module( original_layer.module, original_layer.name, rank, svd_lib_ref) elif isinstance(original_layer.module, nn.Linear): module_a, module_b = WeightSvdModuleSplitter.split_fc_module( original_layer.module, original_layer.name, rank, svd_lib_ref) else: raise TypeError("Only Conv and FC layers are currently supported") # Create a sequential of the split layers seq = nn.Sequential(module_a, module_b) # layer_attr of split layers layer_a = lad.Layer(module_a, original_layer.name + '_a', original_layer.output_shape) layer_b = lad.Layer(module_b, original_layer.name + '_b', original_layer.output_shape) return seq, layer_a, layer_b
def _copy_model(cls, model, model_layers): """ Creates a copy of the original model and its layers :param model: the original model :param model_layers: original model's layers :return: """ # Create a deep-copy of the model to return model_copy = copy.deepcopy(model) # Create an empty model_layers to fill model_layers_copy = {} modules_in_copy = list(model_copy.modules()) # For all modules in the current model for index, module in enumerate(model.modules()): # If this module is included in the existing model_layers, we need to add a corresponding entry into # model_layers_copy if id(module) in model_layers: existing_layer = model_layers[id(module)] new_layer = lad.Layer(modules_in_copy[index], existing_layer.name, existing_layer.output_shape) new_layer.picked_for_compression = existing_layer.picked_for_compression model_layers_copy[id(modules_in_copy[index])] = new_layer # Now we need to set parent references lad.LayerDatabase.set_reference_to_parent_module( model_copy, model_layers_copy) return model_copy, model_layers_copy
def test_calculate_spatial_svd_cost_with_stride(self): conv = nn.Conv2d(32, 64, kernel_size=5, padding=(2, 2), stride=2) layer = lad.Layer(conv, "conv", output_shape=[1, 64, 14, 14]) original_cost = cc.CostCalculator.compute_layer_cost(layer) compressed_cost = cc.SpatialSvdCostCalculator.calculate_cost_given_rank( layer, 40) self.assertEqual(10035200, original_cost.mac) self.assertEqual(5017600, compressed_cost.mac) print(original_cost) print(compressed_cost)
def test_compute_layer_cost(self): logger.debug(self.id()) conv1 = nn.Conv2d(1, 32, kernel_size=5) layer1 = lad.Layer(conv1, "conv1", (1, 32, 28, 28)) cost1 = cc.CostCalculator.compute_layer_cost(layer1) self.assertEqual(32 * 1 * 5 * 5, cost1.memory) self.assertEqual(32 * 1 * 5 * 5 * 28 * 28, cost1.mac) conv2 = nn.Conv2d(32, 64, kernel_size=5) layer2 = lad.Layer(conv2, "conv2", (1, 64, 14, 14)) cost2 = cc.CostCalculator.compute_layer_cost(layer2) self.assertEqual(64 * 32 * 5 * 5, cost2.memory) self.assertEqual(64 * 32 * 5 * 5 * 14 * 14, cost2.mac) conv2 = nn.Conv2d(32, 32, kernel_size=5, groups=32) layer2 = lad.Layer(conv2, "conv2", (1, 32, 14, 14)) cost2 = cc.CostCalculator.compute_layer_cost(layer2) self.assertEqual(32 * 1 * 5 * 5, cost2.memory) self.assertEqual(32 * 1 * 5 * 5 * 14 * 14, cost2.mac)
def test_calculate_weight_svd_cost(self): conv = nn.Conv2d(32, 64, kernel_size=5, padding=(2, 2)) layer = lad.Layer(conv, "conv", output_shape=[1, 64, 28, 28]) self.assertEqual(32, cc.WeightSvdCostCalculator.calculate_max_rank(layer)) comp_ratios_to_check = [0.8, 0.75, 0.5, 0.25, 0.125] original_cost = cc.CostCalculator.compute_layer_cost(layer) for comp_ratio in comp_ratios_to_check: rank = cc.WeightSvdCostCalculator.calculate_rank_given_comp_ratio( layer, comp_ratio, CostMetric.mac) print('Rank = {}, for compression_ratio={}'.format( rank, comp_ratio)) compressed_cost = cc.WeightSvdCostCalculator.calculate_cost_given_rank( layer, rank) print('Compressed cost={}, compression_ratio={}'.format( compressed_cost, compressed_cost.mac / original_cost.mac)) self.assertTrue( math.isclose(compressed_cost.mac / original_cost.mac, comp_ratio, abs_tol=0.03)) # Higher level API for comp_ratio in comp_ratios_to_check: compressed_cost = cc.WeightSvdCostCalculator.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.03))