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))