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