def test_prune_layer(self): orig_model = mnist_torch_model.Net() orig_model.eval() # Create a layer database orig_layer_db = LayerDatabase(orig_model, input_shape=(1, 1, 28, 28)) # Copy the db comp_layer_db = copy.deepcopy(orig_layer_db) 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) input_channel_pruner = InputChannelPruner( data_loader=data_loader, input_shape=(1, 1, 28, 28), num_reconstruction_samples=num_reconstruction_samples, allow_custom_downsample_ops=True) conv2 = comp_layer_db.find_layer_by_name('conv2') input_channel_pruner._prune_layer(orig_layer_db, comp_layer_db, conv2, 0.5, CostMetric.mac) self.assertTrue(comp_layer_db.model.conv2.in_channels, 16) self.assertTrue(comp_layer_db.model.conv2.out_channels, 64)
def test_get_quantized_weight(self): model = mnist_model.Net() params = qsim.QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme=QuantScheme.post_training_tf) use_cuda = False dataset_size = 2 batch_size = 1 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size) def pass_data_through_model(model, early_stopping_iterations=None, use_cuda=False): # forward pass for given number of batches for model for _, (images_in_one_batch, _) in enumerate(data_loader): model(images_in_one_batch) quantsim = qsim.QuantizationSimModel(model=model, quant_scheme=params.quant_scheme, rounding_mode=params.round_mode, default_output_bw=params.act_bw, default_param_bw=params.weight_bw, in_place=False, dummy_input=torch.rand( 1, 1, 28, 28)) quantsim.compute_encodings(pass_data_through_model, None) layer = quantsim.model.conv2 quant_dequant_weights = bias_correction.get_quantized_dequantized_weight( layer, use_cuda) self.assertEqual(quant_dequant_weights.shape, torch.Size([64, 32, 5, 5]))
def bias_correction_analytical_and_empirical(): dataset_size = 2000 batch_size = 64 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(3, 224, 224)) model = MobileNetV2() model.eval() # Find all BN + Conv pairs for analytical BC and remaining Conv for Empirical BC module_prop_dict = bias_correction.find_all_conv_bn_with_activation( model, input_shape=(1, 3, 224, 224)) params = QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme='tf_enhanced') # Perform Bias Correction bias_correction.correct_bias(model.to(device="cuda"), params, num_quant_samples=1000, data_loader=data_loader, num_bias_correct_samples=512, conv_bn_dict=module_prop_dict, perform_only_empirical_bias_corr=False)
def test_get_output_of_layer(self): model = TestNet() dataset_size = 2 batch_size = 2 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size) for images_in_one_batch, _ in data_loader: conv2_output_data = bias_correction.get_output_data( model.conv2, model, images_in_one_batch) # max out number of batches number_of_batches = 1 iterator = data_loader.__iter__() for batch in range(number_of_batches): images_in_one_batch, _ = iterator.__next__() conv1_output = model.conv1(images_in_one_batch) conv2_input = conv1_output conv2_output = model.conv2( functional.relu(functional.max_pool2d(conv2_input, 2))) # compare the output from conv2 layer self.assertTrue( np.allclose( to_numpy(conv2_output), np.asarray(conv2_output_data)[batch * batch_size:(batch + 1) * batch_size, :, :, :]))
def test_bias_correction_empirical_with_config_file(self): # Using a dummy extension of MNIST torch.manual_seed(10) model = mnist_model.Net() model = model.eval() model_copy = copy.deepcopy(model) dataset_size = 2 batch_size = 1 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(1, 28, 28)) # Takes default config file params = qsim.QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme=QuantScheme.post_training_tf, config_file=None) with unittest.mock.patch( 'aimet_torch.bias_correction.call_empirical_mo_correct_bias' ) as empirical_mock: bias_correction.correct_bias(model, params, 2, data_loader, 2) self.assertEqual(empirical_mock.call_count, 4) self.assertTrue( np.allclose(model.conv1.bias.detach().cpu().numpy(), model_copy.conv1.bias.detach().cpu().numpy())) self.assertTrue(model.conv2.bias.detach().cpu().numpy() is not None) self.assertTrue(model.fc1.bias.detach().cpu().numpy() is not None)
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_subsampled_input_data(self, np_choice_function): """ Test to collect activations (input from model_copy and output from model for conv2 layer) and compare with sub sampled input data """ # hardcoded mocked 10 sample locations # (0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (6, 6), (5, 5) np_choice_function.return_value = heights = widths = [ 0, 1, 2, 3, 4, 5, 6, 7, 6, 5 ] orig_model = TestNet() comp_model = copy.deepcopy(orig_model) # only one image and from that 10 samples dataset_size = 1 batch_size = 1 num_reconstruction_samples = 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, image_size=(1, 28, 28)) conv2_input_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) # collect the input data of conv2 from compressed model using same data loader iterator = data_loader.__iter__() images_in_one_batch, _ = iterator.__next__() conv1_output = comp_model.conv1(images_in_one_batch) conv2_input = functional.relu(functional.max_pool2d(conv1_output, 2)) kernel_size_h, kernel_size_w = comp_model.conv2.kernel_size for sample in range(num_reconstruction_samples): self.assertTrue( np.array_equal( conv2_input_data[sample, :, :, :], conv2_input[0, :, heights[sample]:heights[sample] + kernel_size_h, widths[sample]:widths[sample] + kernel_size_w].detach().cpu().numpy()))
def test_get_activation_data(self): """ Test to collect activations (input from model_copy and output from model for conv2 layer) and compare """ orig_model = TestNet().cuda() comp_model = copy.deepcopy(orig_model) dataset_size = 1000 batch_size = 10 # max out number of batches number_of_batches = 100 # 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_input_data, conv2_output_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=number_of_batches) iterator = data_loader.__iter__() for batch in range(number_of_batches): images_in_one_batch, _ = iterator.__next__() conv1_output = orig_model.conv1(images_in_one_batch.cuda()) conv2_input = conv1_output conv2_output = orig_model.conv2( functional.relu(functional.max_pool2d(conv2_input, 2))) # compare the output from conv2 layer self.assertTrue( np.array_equal( to_numpy(conv2_output), conv2_output_data[batch * batch_size:(batch + 1) * batch_size, :, :, :])) conv1_output_copy = comp_model.conv1(images_in_one_batch.cuda()) conv2_input_copy = functional.relu( functional.max_pool2d(conv1_output_copy, 2)) # compare the inputs of conv2 layer self.assertTrue( np.array_equal( to_numpy(conv2_input_copy), conv2_input_data[batch * batch_size:(batch + 1) * batch_size, :, :, :]))
def test_data_sub_sampling_and_reconstruction(self): """Test end to end data sub sampling and reconstruction for MNIST conv2 layer""" orig_model = mnist_model() 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) cp = InputChannelPruner( data_loader=data_loader, input_shape=None, num_reconstruction_samples=num_reconstruction_samples, allow_custom_downsample_ops=True) cp._data_subsample_and_reconstruction(orig_layer=orig_model.conv2, pruned_layer=comp_model.conv2, orig_model=orig_model, comp_model=comp_model) self.assertEqual(comp_model.conv2.weight.data.shape, orig_model.conv2.weight.data.shape) self.assertEqual(comp_model.conv2.bias.data.shape, orig_model.conv2.bias.data.shape) # 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(comp_model.conv2.weight.data), to_numpy(orig_model.conv2.weight.data), atol=1e-0)) self.assertTrue( np.allclose(to_numpy(comp_model.conv2.bias.data), to_numpy(orig_model.conv2.bias.data), atol=1e-0))
def test_subsampled_input_data_fc(self): """ Test to collect activations (input from model_copy and output from model for fc1 layer) and compare with sub sampled output data """ orig_model = TestNet() comp_model = copy.deepcopy(orig_model) # only one image and from that 10 samples dataset_size = 100 batch_size = 10 num_reconstruction_samples = 5000 # create fake data loader with image size (1, 28, 28) data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(1, 28, 28)) fc1_input_data, _ = DataSubSampler.get_sub_sampled_data( orig_layer=orig_model.fc1, pruned_layer=comp_model.fc1, orig_model=orig_model, comp_model=comp_model, data_loader=data_loader, num_reconstruction_samples=num_reconstruction_samples) self.assertTrue(fc1_input_data.shape[0] * fc1_input_data.shape[1] > num_reconstruction_samples) # collect the output data of fc1 from original model using same data loader iterator = data_loader.__iter__() images_in_one_batch, _ = iterator.__next__() conv1_output = orig_model.conv1(images_in_one_batch) conv2_input = conv1_output conv2_output = orig_model.conv2( functional.relu(functional.max_pool2d(conv2_input, 2))) fc1_input = conv2_output fc1_input = functional.relu(functional.max_pool2d(fc1_input, 2)) fc1_input = fc1_input.view(fc1_input.size(0), -1).detach().cpu().numpy() # compare data of first batch only self.assertTrue(np.array_equal(fc1_input_data[0:10], fc1_input))
def test_bias_correction_analytical_and_empirical_ignore_layer(self): torch.manual_seed(10) model = MockMobileNetV1() model = model.eval() dataset_size = 2 batch_size = 1 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(3, 224, 224)) params = qsim.QuantParams(weight_bw=8, act_bw=8, round_mode="nearest", quant_scheme=QuantScheme.post_training_tf) conv_bn_dict = find_all_conv_bn_with_activation(model, input_shape=(1, 3, 224, 224)) layer = model.model[0][0] layers_to_ignore = [layer] with unittest.mock.patch( 'aimet_torch.bias_correction.call_analytical_mo_correct_bias' ) as analytical_mock: with unittest.mock.patch( 'aimet_torch.bias_correction.call_empirical_mo_correct_bias' ) as empirical_mock: bias_correction.correct_bias( model, params, 2, data_loader, 2, conv_bn_dict, perform_only_empirical_bias_corr=False, layers_to_ignore=layers_to_ignore) self.assertEqual(analytical_mock.call_count, 8) # one layer ignored self.assertEqual(empirical_mock.call_count, 9) self.assertTrue( model.model[1][0].bias.detach().cpu().numpy() is not None)
def test_cached_dataset(self): """ Test cache data loader splitting into train and validation """ dataset_size = 256 batch_size = 16 # create fake data loader with image size (1, 2, 2) data_loader = utils.create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(1, 2, 2)) num_batches = 6 path = '/tmp/test_cached_dataset/' cached_dataset = utils.CachedDataset(data_loader, num_batches, path) self.assertEqual(len(cached_dataset), 6) # Try creating cached data loader by more than possible batches from data loader and expect ValueError possible_batches = int(dataset_size / batch_size) with pytest.raises(ValueError): utils.CachedDataset(data_loader, possible_batches + 1, path) shutil.rmtree('/tmp/test_cached_dataset/')
def test_data_sub_sample_and_reconstruction_with_zero_channels(self): """Test end to end data sub sampling and reconstruction for MNIST conv2 layer""" orig_model = mnist_model() 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) cp = InputChannelPruner( data_loader=data_loader, input_shape=None, num_reconstruction_samples=num_reconstruction_samples, allow_custom_downsample_ops=True) input_channels_to_prune = [0, 1, 2, 3, 4, 5, 15, 29, 24, 28] zero_out_input_channels(comp_model.conv2, input_channels_to_prune) before_reconstruction = comp_model.conv2.weight.data cp._data_subsample_and_reconstruction(orig_layer=orig_model.conv2, pruned_layer=comp_model.conv2, orig_model=orig_model, comp_model=comp_model) after_reconstruction = comp_model.conv2.weight.data self.assertEqual(comp_model.conv2.weight.data.shape, orig_model.conv2.weight.data.shape) self.assertEqual(comp_model.conv2.bias.data.shape, orig_model.conv2.bias.data.shape) # make sure they are not same self.assertFalse( np.allclose(to_numpy(before_reconstruction), to_numpy(after_reconstruction)))
def bias_correction_empirical(): dataset_size = 2000 batch_size = 64 data_loader = create_fake_data_loader(dataset_size=dataset_size, batch_size=batch_size, image_size=(3, 224, 224)) model = MobileNetV2() model.eval() params = QuantParams(weight_bw=4, act_bw=4, round_mode="nearest", quant_scheme='tf_enhanced') # Perform Bias Correction bias_correction.correct_bias(model.to(device="cuda"), params, num_quant_samples=1000, data_loader=data_loader.train_loader, num_bias_correct_samples=512)
def test_prune_layer_with_seq(self): """ Test end to end prune layer 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) # Create a layer database orig_layer_db = LayerDatabase(resnet18_model, input_shape=(1, 3, 224, 224)) # Copy the db comp_layer_db = copy.deepcopy(orig_layer_db) 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) conv_below_split = comp_layer_db.find_layer_by_name('layer1.1.conv1') input_channel_pruner._prune_layer(orig_layer_db, comp_layer_db, conv_below_split, 0.25, CostMetric.mac) # 64 * 0.25 = 16 self.assertEqual(comp_layer_db.model.layer1[1].conv1[1].in_channels, 16) 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, 16, 3, 3])
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)