def generate_and_add_module_winnow_list_with_names(
model: torch.nn.Module,
list_of_modules_to_winnow: List[Tuple[torch.nn.Module, List[int]]]):
"""
Generates the module names for the modules to winnow.
Creates a Tuple with Module, list of channels to winnow and Module Name
and adds the Tuple to a list for later use.
:param model: model for which to compare modules from
:param list_of_modules_to_winnow: List of tuples of modules to winnow and corresponding channels to winnow
:return: list of module information.
"""
list_of_module_info = []
if list_of_modules_to_winnow is not None:
model_name = type(model).__name__
logger.debug("Model name: %s", model_name)
for module, list_of_channels_to_winnow, in list_of_modules_to_winnow:
name = get_layer_name(model, module)
# This name doesn't contain the model's name.
# Prepend the model's name to the module's name.
name = '.'.join([model_name, name])
mod_tuple = (module, list_of_channels_to_winnow, name)
list_of_module_info.append(mod_tuple)
return list_of_module_info
def test_find_layer_groups_to_scale_for_network_with_residuals(self):
torch.manual_seed(10)
model = MockMobileNetV2()
model.eval()
fold_all_batch_norms(model, (1, 3, 224, 224))
graph_search = GraphSearchUtils(model, (1, 3, 224, 224))
layer_groups = graph_search.find_layer_groups_to_scale()
self.assertEqual(4, len(layer_groups))
self.assertIn([
model.features[3].conv[0], model.features[3].conv[3],
model.features[3].conv[6]
], layer_groups)
self.assertIn([
model.features[4].conv[0], model.features[4].conv[3],
model.features[4].conv[6]
], layer_groups)
self.assertIn([
model.features[5].conv[0], model.features[5].conv[3],
model.features[5].conv[6], model.features[6][0]
], layer_groups)
for layer_group in layer_groups:
print("Group ------- ")
for module in layer_group:
print(" " + get_layer_name(model, module))
def test_mask_propagation_through_concat(self):
""" After the graph is constructed, the Op should have default masks and connectivity for all module types. """
logger.debug("Test default mask and connectivity.")
model = SingleConcat()
# Test forward pass on the copied model before zeroing out channels of layers.
input_shape = [1, 3, 224, 224]
module_zero_channels_list = []
module = model.conv4
input_channels_to_prune = [
1, 3, 5, 7, 9, 19, 21, 23, 25, 43, 45, 47, 49, 51, 57, 59, 61, 63
]
module_mask_pair = (module, input_channels_to_prune)
module_zero_channels_list.append(module_mask_pair)
print("Order of modules in in the API:",
[get_layer_name(model, m) for m, _ in module_zero_channels_list])
# API version 2.
winnowed_model, _ = winnow_model(model,
input_shape,
module_zero_channels_list,
in_place=True,
verbose=True)
# validate winnowed net
# input_tensor = torch.rand(input_shape).double()
input_tensor = torch.rand(input_shape)
# winnowed_model.double()
winnowed_model.eval()
_ = winnowed_model(input_tensor)
self.assertEqual(0, 0)
def test_data_sub_sampling_and_reconstruction_without_bias(self):
"""Test end to end data sub sampling and reconstruction for MNIST conv2 layer (without bias)"""
orig_model = mnist_model()
# set bias to None
orig_model.conv2.bias = None
comp_model = copy.deepcopy(orig_model)
dataset_size = 100
batch_size = 10
# max out number of batches
number_of_batches = 10
samples_per_image = 10
num_reconstruction_samples = number_of_batches * batch_size * samples_per_image
# create fake data loader with image size (1, 28, 28)
data_loader = create_fake_data_loader(dataset_size=dataset_size,
batch_size=batch_size)
conv2_pr_layer_name = get_layer_name(comp_model, comp_model.conv2)
sampled_inp_data, sampled_out_data = DataSubSampler.get_sub_sampled_data(
orig_layer=orig_model.conv2,
pruned_layer=comp_model.conv2,
orig_model=orig_model,
comp_model=comp_model,
data_loader=data_loader,
num_reconstruction_samples=num_reconstruction_samples)
conv_layer = get_layer_by_name(model=comp_model,
layer_name=conv2_pr_layer_name)
assert conv_layer == comp_model.conv2
# original weight before reconstruction
orig_weight = conv_layer.weight.data
WeightReconstructor.reconstruct_params_for_conv2d(
layer=conv_layer,
input_data=sampled_inp_data,
output_data=sampled_out_data)
# new weight after reconstruction
new_weight = conv_layer.weight.data
new_bias = conv_layer.bias
self.assertEqual(new_weight.shape, orig_weight.shape)
self.assertEqual(new_bias, None)
# if you increase the data (data set size, number of batches or samples per image),
# reduce the absolute tolerance
self.assertTrue(
np.allclose(to_numpy(new_weight), to_numpy(orig_weight),
atol=1e-0))
def test_mask_propagation_through_single_chunk(self):
""" After the graph is constructed, the Op should have default masks and connectivity for all module types. """
logger.debug("Test default mask and connectivity.")
model = SingleChunk()
# Test forward pass on the copied model before zering out channels of layers.
input_shape = [1, 3, 224, 224]
module_zero_channels_list = []
module = model.conv4
input_channels_to_prune = [5, 9]
module_mask_pair = (module, input_channels_to_prune)
module_zero_channels_list.append(module_mask_pair)
print("Order of modules in in the API:",
[get_layer_name(model, m) for m, _ in module_zero_channels_list])
# API version 2.
_, _ = winnow_model(model,
input_shape,
module_zero_channels_list,
in_place=True,
verbose=True)
self.assertEqual(0, 0)
def test_conv_to_conv_mask_propagation(self):
""" After the graph is constructed, the Op should have default masks and connectivity for all module types. """
logger.debug("Test default mask and connectivity.")
model = SingleResidual()
# Test forward pass on the copied model before winnowing.
input_shape = [1, 3, 224, 224]
input_tensor = torch.rand(input_shape).double()
model.double()
model.eval()
print(
"test_conv_to_conv_mask_propagation(): Testing forward pass before winnowing."
)
validation_output = model(input_tensor)
# Convert the model back to float.
model.float()
module_zero_channels_list = []
module = model.conv3
input_channels_to_prune = [1, 3]
module_mask_pair = (module, input_channels_to_prune)
module_zero_channels_list.append(module_mask_pair)
print("Order of modules in in the API:",
[get_layer_name(model, m) for m, _ in module_zero_channels_list])
# API version 2.
winnowed_model, _ = winnow_model(model,
input_shape,
module_zero_channels_list,
in_place=True,
verbose=True)
# validate winnowed net
input_tensor = torch.rand(input_shape).double()
winnowed_model.double()
winnowed_model.eval()
test_output = winnowed_model(input_tensor)
self.assertTrue(test_output.shape == validation_output.shape)
# self.assertTrue(test_output.allclose(validation_output)) # TBD. Why is this failing?
print(test_output)
print(validation_output)
# In the winnowed model, conv3 has in_channels = 62, out_channels = 64
self.assertTrue(winnowed_model.conv3.in_channels == 62)
self.assertTrue(winnowed_model.conv3.out_channels == 64)
# The winnowed model's bn2 layer has 62 num_features
self.assertEqual(winnowed_model.bn2.num_features, 62)
self.assertEqual(list(winnowed_model.bn2.weight.shape), [62])
self.assertEqual(list(winnowed_model.bn2.bias.shape), [62])
self.assertEqual(list(winnowed_model.bn2.running_mean.shape), [62])
self.assertEqual(list(winnowed_model.bn2.running_var.shape), [62])
# In the winnowed model, conv2 has in_channels = 64, out_channels = 62 (impacted by layer3 pruning)
self.assertTrue(winnowed_model.conv2.in_channels == 64)
self.assertTrue(winnowed_model.conv2.out_channels == 62)
print(
"test_conv_to_conv_mask_propagation(): Successfully validated winnowed model."
)