def test_sort_on_occurrence(self):
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=3)
self.conv2 = nn.Conv2d(10, 10, kernel_size=3)
self.conv3 = nn.Conv2d(10, 10, kernel_size=3)
self.conv4 = nn.Conv2d(10, 10, kernel_size=3)
self.fc1 = nn.Linear(490, 300)
self.fc2 = nn.Linear(300, 10)
def forward(self, x):
x = functional.relu(functional.max_pool2d(self.conv1(x), 2))
x = functional.relu(self.conv2(x))
x = functional.relu(self.conv3(x))
x = functional.relu(self.conv4(x))
x = x.view(x.size(0), -1)
x = functional.relu(self.fc1(x))
x = self.fc2(x)
return functional.log_softmax(x, dim=1)
orig_model = Net()
data_loader = unittest.mock.MagicMock()
number_of_batches = unittest.mock.MagicMock()
samples_per_image = unittest.mock.MagicMock()
input_channel_pruner = InputChannelPruner(
data_loader=data_loader,
input_shape=(1, 1, 28, 28),
num_reconstruction_samples=number_of_batches,
allow_custom_downsample_ops=True)
layer_comp_ratio_list = [
LayerCompRatioPair(Layer(orig_model.conv4, None, None), None),
LayerCompRatioPair(Layer(orig_model.conv1, None, None), None),
LayerCompRatioPair(Layer(orig_model.conv3, None, None), None),
LayerCompRatioPair(Layer(orig_model.conv2, None, None), None)
]
sorted_layer_comp_ratio_list = input_channel_pruner._sort_on_occurrence(
orig_model, layer_comp_ratio_list)
self.assertEqual(sorted_layer_comp_ratio_list[0].layer.module,
orig_model.conv1)
self.assertEqual(sorted_layer_comp_ratio_list[1].layer.module,
orig_model.conv2)
self.assertEqual(sorted_layer_comp_ratio_list[2].layer.module,
orig_model.conv3)
self.assertEqual(sorted_layer_comp_ratio_list[3].layer.module,
orig_model.conv4)
self.assertTrue(
isinstance(sorted_layer_comp_ratio_list[0].layer, Layer))
self.assertTrue(
isinstance(sorted_layer_comp_ratio_list[1].layer, Layer))
self.assertTrue(
isinstance(sorted_layer_comp_ratio_list[2].layer, Layer))
self.assertTrue(
isinstance(sorted_layer_comp_ratio_list[3].layer, Layer))
def _update_layer_database_after_winnowing(comp_layer_db,
updated_module_list):
# Get reference to the new module to update in the layer database
for old_module_name, new_module in updated_module_list:
try:
old_layer = comp_layer_db.find_layer_by_name(old_module_name)
except KeyError:
# Nothing to update, LayerDatabase was not tracking this layer
continue
if isinstance(new_module, torch.nn.Sequential):
comp_layer_db.update_layer_with_module_in_sequential(
old_layer, new_module)
else:
# Determine new output shape
new_output_shape = [
old_layer.output_shape[0], new_module.out_channels,
old_layer.output_shape[2], old_layer.output_shape[3]
]
new_layer = Layer(new_module, old_layer.name, new_output_shape)
comp_layer_db.replace_layer(old_layer, new_layer)
def test_split_fc_layer_without_mo(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
with unittest.mock.patch('aimet_torch.layer_database.LayerDatabase'):
with unittest.mock.patch('aimet_torch.svd.layer_selector_deprecated.LayerSelectorDeprecated'):
svd = s.SvdImpl(model=model, run_model=None, run_model_iterations=1, input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd, cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.memory,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.top_n_layers, num_layers=2)
layer_attr = Layer(model.fc1, id(model.fc1), [3136, 1024, 1, 1])
svd._svd_lib_ref = create_autospec(pymo.Svd, instance=True)
split_weights = [np.zeros((400, model.fc1.in_features)).flatten().tolist(),
np.zeros((model.fc1.out_features, 400)).flatten().tolist()]
svd._svd_lib_ref.SplitLayerWeights.return_value = split_weights
split_biases = [np.zeros(400).flatten().tolist(),
np.zeros(model.fc1.out_features).flatten().tolist()]
svd._svd_lib_ref.SplitLayerBiases.return_value = split_biases
split_layer = svd_pruner_deprecated.DeprecatedSvdPruner
seq, layer_a_attr, layer_b_attr = split_layer.prune_layer(layer_attr, 400, svd_lib_ref=svd._svd_lib_ref)
self.assertEqual((400, model.fc1.in_features), seq[0].weight.shape)
self.assertEqual([400], list(seq[0].bias.shape))
self.assertEqual((model.fc1.out_features, 400), seq[1].weight.shape)
self.assertEqual([model.fc1.out_features], list(seq[1].bias.shape))
self.assertEqual(layer_a_attr.module, seq[0])
self.assertEqual(layer_b_attr.module, seq[1])
def test_find_layer_comp_ratio_given_eval_score(self):
eval_scores_dict = {
'layer1': {
Decimal('0.1'): 90,
Decimal('0.5'): 50,
Decimal('0.7'): 30,
Decimal('0.8'): 20
},
'layer2': {
Decimal('0.1'): 11,
Decimal('0.3'): 23,
Decimal('0.5'): 47,
Decimal('0.7'): 85,
Decimal('0.9'): 89
}
}
layer2 = Layer(nn.Conv2d(32, 64, 3), "layer2", None)
greedy_algo = comp_ratio_select.GreedyCompRatioSelectAlgo
comp_ratio = greedy_algo._find_layer_comp_ratio_given_eval_score(
eval_scores_dict, 45, layer2)
self.assertEqual(Decimal('0.5'), comp_ratio)
comp_ratio = greedy_algo._find_layer_comp_ratio_given_eval_score(
eval_scores_dict, 48, layer2)
self.assertEqual(Decimal('0.7'), comp_ratio)
comp_ratio = greedy_algo._find_layer_comp_ratio_given_eval_score(
eval_scores_dict, 90, layer2)
self.assertEqual(None, comp_ratio)
def test_set_parent_attribute_two_deep(self):
"""With a two-deep model"""
class SubNet(nn.Module):
def __init__(self):
super(SubNet, self).__init__()
self.conv1 = nn.Conv2d(30, 40, 5)
self.conv2 = nn.Conv2d(40, 50, kernel_size=5)
def forward(self, *inputs):
pass
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, 5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.subnet1 = SubNet()
self.fc1 = nn.Linear(320, 50)
self.subnet2 = SubNet()
self.fc2 = nn.Linear(50, 10)
def forward(self, *inputs):
pass
net = Net()
model = net.to("cpu")
# create layer attribute
output_activation_shape = None
layers = {id(model.subnet1.conv2): Layer(model.subnet1.conv2, id(model.subnet1.conv2),
output_activation_shape),
id(model.subnet2.conv1): Layer(model.subnet2.conv1, id(model.subnet2.conv1),
output_activation_shape),
id(model.fc2): Layer(model.fc2, id(model.fc2),
output_activation_shape)}
LayerDatabase.set_reference_to_parent_module(model, layers)
# child : model.subnet1.conv2 --> parent : model.subnet1
self.assertEqual(model.subnet1, layers[id(model.subnet1.conv2)].parent_module)
# child : model.subnet2.conv1 --> parent : model.subnet2
self.assertEqual(model.subnet2, layers[id(model.subnet2.conv1)].parent_module)
# child : model.fc2 --> parent : model
self.assertEqual(model, layers[id(model.fc2)].parent_module)
def test_select_all_conv_layers(self):
mock_output_shape = (1, 1, 1, 1)
# Two regular conv layers
layer1 = Layer(Conv2d(10, 20, 5), '', mock_output_shape)
layer2 = Layer(Conv2d(10, 20, 5), '', mock_output_shape)
layer3 = Layer(Conv2d(10, 10, 5, groups=10), '', mock_output_shape)
layer_db = MagicMock()
layer_db.__iter__.return_value = [layer1, layer2, layer3]
layer_selector = ConvNoDepthwiseLayerSelector()
layer_selector.select(layer_db, [])
layer_db.mark_picked_layers.assert_called_once_with([layer1, layer2])
# One conv and one linear layer
layer1 = Layer(Conv2d(10, 20, 5), '', mock_output_shape)
layer2 = Layer(Linear(10, 20), '', mock_output_shape)
layer_db = MagicMock()
layer_db.__iter__.return_value = [layer1, layer2]
layer_selector.select(layer_db, [])
layer_db.mark_picked_layers.assert_called_once_with([layer1])
# Two regular conv layers - one in ignore list
layer1 = Layer(Conv2d(10, 20, 5), '', mock_output_shape)
layer2 = Layer(Conv2d(10, 20, 5), '', mock_output_shape)
layer_db = MagicMock()
layer_db.__iter__.return_value = [layer1, layer2]
layer_selector.select(layer_db, [layer2.module])
layer_db.mark_picked_layers.assert_called_once_with([layer1])
def test_set_attributes_with_sequentials(self):
"""With a one-deep model"""
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, 5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.subnet1 = nn.Sequential(
nn.Conv2d(1, 10, 5),
nn.ReLU(),
nn.Conv2d(10, 20, 5)
)
self.fc1 = nn.Linear(320, 50)
self.subnet2 = nn.Sequential(
nn.Conv2d(1, 10, 5),
nn.ReLU(),
nn.Conv2d(1, 10, 5)
)
self.fc2 = nn.Linear(50, 10)
def forward(self, *inputs):
pass
net = Net()
model = net.to("cpu")
# create layer attribute
output_activation_shape = None
layers = {id(model.subnet1[2]): Layer(model.subnet1[2], id(model.subnet1[2]),
output_activation_shape),
id(model.subnet2[0]): Layer(model.subnet2[0], id(model.subnet2[0]),
output_activation_shape),
id(model.fc2): Layer(model.fc2, id(model.fc2),
output_activation_shape)}
LayerDatabase.set_reference_to_parent_module(model, layers)
# child : model.subnet1.2 --> parent : model.subnet1
self.assertEqual(model.subnet1, layers[id(model.subnet1[2])].parent_module)
# child : model.subnet2.1 --> parent : model.subnet2
self.assertEqual(model.subnet2, layers[id(model.subnet2[0])].parent_module)
# child : model.fc2 --> parent : model
self.assertEqual(model, layers[id(model.fc2)].parent_module)
def test_create_compressed_model(self):
AimetLogger.set_level_for_all_areas(logging.DEBUG)
logger.debug(self.id())
model = MnistModel().to("cpu")
with unittest.mock.patch('aimet_torch.svd.layer_database.LayerDatabase'):
with unittest.mock.patch('aimet_torch.svd.layer_selector_deprecated.LayerSelectorDeprecated'):
svd = s.SvdImpl(model=model, run_model=None, run_model_iterations=1, input_shape=(1, 1, 28, 28),
compression_type=aimet_torch.svd.svd_intf_defs_deprecated.CompressionTechnique.svd, cost_metric=aimet_torch.svd.svd_intf_defs_deprecated.CostMetric.memory,
layer_selection_scheme=aimet_torch.svd.svd_intf_defs_deprecated.LayerSelectionScheme.top_n_layers, num_layers=2)
ls.LayerSelectorDeprecated._pick_compression_layers = create_autospec(ls.LayerSelectorDeprecated._pick_compression_layers)
layer_attr1 = Layer(model.fc2, id(model.fc2), model.fc2.weight.shape)
layer_attr1.parent_module = model
layer_attr1.var_name_of_module_in_parent = "fc2"
layer_attr1.output_shape = [0, 0, 1, 1]
layer_attr1.name = 'fc2'
layer_attr2 = Layer(model.conv2, id(model.conv2), model.conv2.weight.shape)
layer_attr2.parent_module = model
layer_attr2.var_name_of_module_in_parent = "conv2"
layer_attr2.name = 'conv2'
layer_attr1.output_shape = [0, 0, 14, 14]
ls.LayerSelectorDeprecated._pick_compression_layers.return_value = [layer_attr1, layer_attr2]
svd._compressible_layers = {id(model.conv2): layer_attr2,
id(model.fc2): layer_attr1}
ls.LayerSelectorDeprecated._perform_layer_selection(model)
svd._select_candidate_ranks(20)
svd_rank_pair_dict = {'conv2': (31,0), 'fc2': (9,0)}
c_model, c_layer_attr, _ = svd._create_compressed_model(svd_rank_pair_dict)
self.assertTrue(c_model is not model)
self.assertTrue(c_model.conv1 is not model.conv1)
self.assertTrue(c_model.conv2 is not model.conv2)
self.assertFalse(isinstance(svd._model, nn.Sequential))
self.assertEqual((9, 1024), c_model.fc2[0].weight.shape)
self.assertEqual([9], list(c_model.fc2[0].bias.shape))
self.assertEqual((10, 9), c_model.fc2[1].weight.shape)
self.assertEqual([10], list(c_model.fc2[1].bias.shape))
self.assertEqual((31, 32, 1, 1), c_model.conv2[0].weight.shape)
self.assertEqual([31], list(c_model.conv2[0].bias.shape))
self.assertEqual((64, 31, 5, 5), c_model.conv2[1].weight.shape)
self.assertEqual([64], list(c_model.conv2[1].bias.shape))
self.assertEqual(svd._model.conv1.weight.shape, c_model.conv1.weight.shape)
self.assertEqual(svd._model.fc1.weight.shape, c_model.fc1.weight.shape)
# Expect double the number of layers in layer_attr_list
self.assertEqual(4, len(c_layer_attr))
def test_prune_model_with_seq(self):
"""Test end to end prune model with resnet18"""
batch_size = 2
dataset_size = 1000
number_of_batches = 1
samples_per_image = 10
num_reconstruction_samples = number_of_batches * batch_size * samples_per_image
resnet18_model = models.resnet18(pretrained=True)
resnet18_model.eval()
# Create a layer database
orig_layer_db = LayerDatabase(resnet18_model,
input_shape=(1, 3, 224, 224))
data_loader = create_fake_data_loader(dataset_size=dataset_size,
batch_size=batch_size,
image_size=(3, 224, 224))
input_channel_pruner = InputChannelPruner(
data_loader=data_loader,
input_shape=(1, 3, 224, 224),
num_reconstruction_samples=num_reconstruction_samples,
allow_custom_downsample_ops=True)
# keeping compression ratio = 0.5 for all layers
layer_comp_ratio_list = [
LayerCompRatioPair(
Layer(resnet18_model.layer4[1].conv1, 'layer4.1.conv1', None),
0.5),
LayerCompRatioPair(
Layer(resnet18_model.layer3[1].conv1, 'layer3.1.conv1', None),
0.5),
LayerCompRatioPair(
Layer(resnet18_model.layer2[1].conv1, 'layer2.1.conv1', None),
0.5),
LayerCompRatioPair(
Layer(resnet18_model.layer1[1].conv1, 'layer1.1.conv1', None),
0.5),
LayerCompRatioPair(
Layer(resnet18_model.layer1[0].conv2, 'layer1.0.conv2', None),
0.5)
]
comp_layer_db = input_channel_pruner.prune_model(orig_layer_db,
layer_comp_ratio_list,
CostMetric.mac,
trainer=None)
# 1) not below split
self.assertEqual(comp_layer_db.model.layer1[0].conv2.in_channels, 32)
self.assertEqual(comp_layer_db.model.layer1[0].conv2.out_channels, 64)
self.assertEqual(
list(comp_layer_db.model.layer1[0].conv2.weight.shape),
[64, 32, 3, 3])
# impacted
self.assertEqual(comp_layer_db.model.layer1[0].conv1.in_channels, 64)
self.assertEqual(comp_layer_db.model.layer1[0].conv1.out_channels, 32)
self.assertEqual(
list(comp_layer_db.model.layer1[0].conv1.weight.shape),
[32, 64, 3, 3])
# 2) below split
# 64 * .5
self.assertEqual(comp_layer_db.model.layer1[1].conv1[1].in_channels,
32)
self.assertEqual(comp_layer_db.model.layer1[1].conv1[1].out_channels,
64)
self.assertEqual(
list(comp_layer_db.model.layer1[1].conv1[1].weight.shape),
[64, 32, 3, 3])
# 128 * .5
self.assertEqual(comp_layer_db.model.layer2[1].conv1[1].in_channels,
64)
self.assertEqual(comp_layer_db.model.layer2[1].conv1[1].out_channels,
128)
self.assertEqual(
list(comp_layer_db.model.layer2[1].conv1[1].weight.shape),
[128, 64, 3, 3])
# 256 * .5
self.assertEqual(comp_layer_db.model.layer3[1].conv1[1].in_channels,
128)
self.assertEqual(comp_layer_db.model.layer3[1].conv1[1].out_channels,
256)
self.assertEqual(
list(comp_layer_db.model.layer3[1].conv1[1].weight.shape),
[256, 128, 3, 3])
# 512 * .5
self.assertEqual(comp_layer_db.model.layer4[1].conv1[1].in_channels,
256)
self.assertEqual(comp_layer_db.model.layer4[1].conv1[1].out_channels,
512)
self.assertEqual(
list(comp_layer_db.model.layer4[1].conv1[1].weight.shape),
[512, 256, 3, 3])
def test_prune_model(self):
"""Test end to end prune model with Mnist"""
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=3)
self.max_pool2d = nn.MaxPool2d(2)
self.relu1 = nn.ReLU()
self.conv2 = nn.Conv2d(10, 20, kernel_size=3)
self.relu2 = nn.ReLU()
self.conv3 = nn.Conv2d(20, 30, kernel_size=3)
self.relu3 = nn.ReLU()
self.conv4 = nn.Conv2d(30, 40, kernel_size=3)
self.relu4 = nn.ReLU()
self.fc1 = nn.Linear(7 * 7 * 40, 300)
self.relu5 = nn.ReLU()
self.fc2 = nn.Linear(300, 10)
self.log_softmax = nn.LogSoftmax(dim=1)
def forward(self, x):
x = self.relu1(self.max_pool2d(self.conv1(x)))
x = self.relu2(self.conv2(x))
x = self.relu3(self.conv3(x))
x = self.relu4(self.conv4(x))
x = x.view(x.size(0), -1)
x = self.relu5(self.fc1(x))
x = self.fc2(x)
return self.log_softmax(x)
orig_model = Net()
orig_model.eval()
# Create a layer database
orig_layer_db = LayerDatabase(orig_model, input_shape=(1, 1, 28, 28))
dataset_size = 1000
batch_size = 10
# max out number of batches
number_of_batches = 100
samples_per_image = 10
# create fake data loader with image size (1, 28, 28)
data_loader = create_fake_data_loader(dataset_size=dataset_size,
batch_size=batch_size)
input_channel_pruner = InputChannelPruner(
data_loader=data_loader,
input_shape=(1, 1, 28, 28),
num_reconstruction_samples=number_of_batches,
allow_custom_downsample_ops=True)
# keeping compression ratio = 0.5 for all layers
layer_comp_ratio_list = [
LayerCompRatioPair(Layer(orig_model.conv4, 'conv4', None), 0.5),
LayerCompRatioPair(Layer(orig_model.conv3, 'conv3', None), 0.5),
LayerCompRatioPair(Layer(orig_model.conv2, 'conv2', None), 0.5)
]
comp_layer_db = input_channel_pruner.prune_model(orig_layer_db,
layer_comp_ratio_list,
CostMetric.mac,
trainer=None)
self.assertEqual(comp_layer_db.model.conv2.in_channels, 5)
self.assertEqual(comp_layer_db.model.conv2.out_channels, 10)
self.assertEqual(comp_layer_db.model.conv3.in_channels, 10)
self.assertEqual(comp_layer_db.model.conv3.out_channels, 15)
self.assertEqual(comp_layer_db.model.conv4.in_channels, 15)
self.assertEqual(comp_layer_db.model.conv4.out_channels, 40)