def test_collect_inp_out_data_quantsim_model_gpu(self): """ test collect input output data from module """ device_list = [torch.device('cuda:0')] for device in device_list: model = TinyModel().to(device=device) model_input = torch.randn(1, 3, 32, 32).to(device=device) sim = QuantizationSimModel(model, dummy_input=torch.rand( 1, 3, 32, 32).to(device=device)) module_data = utils.ModuleData(model, model.fc) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=True) fc_out = sim.model(model_input) self.assertFalse( np.array_equal(utils.to_numpy(out), utils.to_numpy(fc_out))) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=True, collect_output=False) self.assertTrue( np.array_equal(utils.to_numpy(inp), utils.to_numpy(model_input)))
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 _hook_to_collect_output_data(module, _, out_data): """ hook to collect output data """ out_data = utils.to_numpy(out_data) orig_layer_out_data.append(out_data) raise StopForwardException
def _hook_to_collect_input_data(module, inp_data, _): """ hook to collect input data """ inp_data = utils.to_numpy(inp_data[0]) pruned_layer_inp_data.append(inp_data) raise StopForwardException
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 _collect_inp_out_data(self, device): model = TinyModel().to(device=device) model_input = torch.randn(1, 3, 32, 32).to(device=device) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=False) self.assertEqual(inp, None) self.assertEqual(out, None) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=True, collect_output=False) self.assertTrue( np.array_equal(utils.to_numpy(inp), utils.to_numpy(model_input))) self.assertEqual(out, None) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=True) conv1_out = model.conv1(model_input) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(conv1_out))) self.assertEqual(inp, None) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=True, collect_output=True) conv1_out = model.conv1(model_input) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(conv1_out))) self.assertTrue( np.array_equal(utils.to_numpy(inp), utils.to_numpy(model_input))) module_data = utils.ModuleData(model, model.fc) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=True) fc_out = model(model_input) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(fc_out))) self.assertEqual(inp, None)
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 _collect_inp_out_data_multi_input(self, device): model = MultiInput().to(device=device) inp_shape_1 = (1, 3, 32, 32) inp_shape_2 = (1, 3, 20, 20) model_input = utils.create_rand_tensors_given_shapes( [inp_shape_1, inp_shape_2]) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=True, collect_output=False) self.assertTrue( np.array_equal(utils.to_numpy(inp), utils.to_numpy(model_input[0]))) self.assertEqual(out, None) module_data = utils.ModuleData(model, model.conv1) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=True) conv1_out = model.conv1(model_input[0]) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(conv1_out))) self.assertEqual(inp, None) module_data = utils.ModuleData(model, model.conv3) inp, out = module_data.collect_inp_out_data(model_input, collect_input=True, collect_output=True) conv3_out = model.conv3(model_input[1]) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(conv3_out))) self.assertTrue( np.array_equal(utils.to_numpy(inp), utils.to_numpy(model_input[1]))) module_data = utils.ModuleData(model, model.fc) inp, out = module_data.collect_inp_out_data(model_input, collect_input=False, collect_output=True) fc_out = model(*model_input) self.assertTrue( np.array_equal(utils.to_numpy(out), utils.to_numpy(fc_out))) self.assertEqual(inp, None)
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_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 test_reconstruct_weight_and_bias_for_layer(self): """ """ model = TestNet() layer = model.conv2 # input shape should be [Ns, Nic, k_h, k_w] number_of_images = 500 inputs = np.random.rand(number_of_images, layer.in_channels, layer.kernel_size[0], layer.kernel_size[1]) outputs = functional.conv2d(torch.FloatTensor(inputs), layer.weight, bias=layer.bias, stride=layer.stride, padding=layer.padding) # expecting output shape (number_of_images, layer.out_channels, 1, 1) self.assertEqual(np.prod(outputs.shape[1:4]), layer.out_channels) reshaped_outputs = outputs.reshape( [outputs.shape[0], np.prod(outputs.shape[1:4])]) WeightReconstructor.reconstruct_params_for_conv2d( layer=layer, input_data=inputs, output_data=to_numpy(reshaped_outputs)) new_outputs = functional.conv2d(torch.FloatTensor(inputs), layer.weight, bias=layer.bias, stride=layer.stride, padding=layer.padding) # if data is increased, choose tolerance wisely self.assertTrue( np.allclose(to_numpy(outputs), to_numpy(new_outputs), atol=1e-5))