diff --git a/test/quantization/core/experimental/test_quantizer.py b/test/quantization/core/experimental/test_quantizer.py index 1311fa3b78a3..d689ee8e99e1 100644 --- a/test/quantization/core/experimental/test_quantizer.py +++ b/test/quantization/core/experimental/test_quantizer.py @@ -184,6 +184,43 @@ def test_dequantize_quantize_rand_b6(self): self.assertTrue(torch.equal(original_input, result)) + r""" Tests for correct dimensions in dequantize_apot result + on random 3-dim tensor with random dimension sizes + and hardcoded values for b, k. + Dequant an input tensor and verify that + dimensions are same as input. + * tensor2quantize: Tensor + * b: 4 + * k: 2 + """ + def test_dequantize_dim(self): + # make observer + observer = APoTObserver(4, 2) + + # generate random size of tensor2quantize between 1 -> 20 + size1 = random.randint(1, 20) + size2 = random.randint(1, 20) + size3 = random.randint(1, 20) + + # make tensor2quantize: random fp values between 0 -> 1000 + tensor2quantize = 1000 * torch.rand(size1, size2, size3, dtype=torch.float) + + observer.forward(tensor2quantize) + + alpha, gamma, quantization_levels, level_indices = observer.calculate_qparams(signed=False) + + # make mock apot_tensor + original_apot = quantize_APoT(tensor2quantize=tensor2quantize, + alpha=alpha, + gamma=gamma, + quantization_levels=quantization_levels, + level_indices=level_indices) + + # dequantize apot_tensor + dequantize_result = dequantize_APoT(apot_tensor=original_apot) + + self.assertEqual(original_apot.data.size(), dequantize_result.size()) + def test_q_apot_alpha(self): with self.assertRaises(NotImplementedError): APoTQuantizer.q_apot_alpha(self) diff --git a/torch/ao/quantization/experimental/quantizer.py b/torch/ao/quantization/experimental/quantizer.py index 46d792db8147..389443530ebb 100644 --- a/torch/ao/quantization/experimental/quantizer.py +++ b/torch/ao/quantization/experimental/quantizer.py @@ -57,15 +57,18 @@ def quantize(self, tensor2quantize: Tensor): result: fp representation of input Tensor """ def dequantize(self, apot_tensor) -> Tensor: - apot_tensor_data = apot_tensor.data + orig_size = apot_tensor.data.size() + apot_tensor_data = apot_tensor.data.flatten() + + print(apot_tensor_data) # map apot_to_float over tensor2quantize elements - result_temp = np.empty(apot_tensor_data.size()) - for ele in apot_tensor_data: - new_ele = apot_to_float(ele, self.quantization_levels, self.level_indices) - np.append(result_temp, new_ele) + result_temp = np.empty(shape=apot_tensor_data.size()) + for i in range(len(apot_tensor_data)): + new_ele = apot_to_float(apot_tensor_data[i], self.quantization_levels, self.level_indices) + result_temp[i] = new_ele - result = torch.from_numpy(result_temp).int() + result = torch.from_numpy(result_temp).reshape(orig_size) return result