# Owner(s): ["oncall: quantization"] from torch.ao.quantization.experimental.observer import APoTObserver import unittest import torch class TestNonUniformObserver(unittest.TestCase): """ Test case 1: calculate_qparams Test that error is thrown when k == 0 """ def test_calculate_qparams_invalid(self): obs = APoTObserver(b=0, k=0) obs.min_val = torch.tensor([0.0]) obs.max_val = torch.tensor([0.0]) with self.assertRaises(AssertionError): alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=False) """ Test case 2: calculate_qparams APoT paper example: https://arxiv.org/pdf/1909.13144.pdf Assume hardcoded parameters: * b = 4 (total number of bits across all terms) * k = 2 (base bitwidth, i.e. bitwidth of every term) * n = 2 (number of additive terms) * note: b = k * n """ def test_calculate_qparams_2terms(self): obs = APoTObserver(b=4, k=2) obs.min_val = torch.tensor([0.0]) obs.max_val = torch.tensor([1.0]) alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=False) alpha_test = torch.max(-obs.min_val, obs.max_val) # check alpha value self.assertEqual(alpha, alpha_test) # calculate expected gamma value gamma_test = 0 for i in range(2): gamma_test += 2**(-i) gamma_test = 1 / gamma_test # check gamma value self.assertEqual(gamma, gamma_test) # check quantization levels size quantlevels_size_test = int(len(quantization_levels)) quantlevels_size = 2**4 self.assertEqual(quantlevels_size_test, quantlevels_size) # check level indices size levelindices_size_test = int(len(level_indices)) self.assertEqual(levelindices_size_test, 16) # check level indices unique values level_indices_test_list = level_indices.tolist() self.assertEqual(len(level_indices_test_list), len(set(level_indices_test_list))) """ Test case 3: calculate_qparams Assume hardcoded parameters: * b = 6 (total number of bits across all terms) * k = 2 (base bitwidth, i.e. bitwidth of every term) * n = 3 (number of additive terms) """ def test_calculate_qparams_3terms(self): obs = APoTObserver(b=6, k=2) obs.min_val = torch.tensor([0.0]) obs.max_val = torch.tensor([1.0]) alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=False) alpha_test = torch.max(-obs.min_val, obs.max_val) # check alpha value self.assertEqual(alpha, alpha_test) # calculate expected gamma value gamma_test = 0 for i in range(3): gamma_test += 2**(-i) gamma_test = 1 / gamma_test # check gamma value self.assertEqual(gamma, gamma_test) # check quantization levels size quantlevels_size_test = int(len(quantization_levels)) quantlevels_size = 2**6 self.assertEqual(quantlevels_size_test, quantlevels_size) # check level indices size levelindices_size_test = int(len(level_indices)) self.assertEqual(levelindices_size_test, 64) # check level indices unique values level_indices_test_list = level_indices.tolist() self.assertEqual(len(level_indices_test_list), len(set(level_indices_test_list))) """ Test case 4: calculate_qparams Same as test case 2 but with signed = True Assume hardcoded parameters: * b = 4 (total number of bits across all terms) * k = 2 (base bitwidth, i.e. bitwidth of every term) * n = 2 (number of additive terms) * signed = True """ def test_calculate_qparams_signed(self): obs = APoTObserver(b=4, k=2) obs.min_val = torch.tensor([0.0]) obs.max_val = torch.tensor([1.0]) alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=True) alpha_test = torch.max(-obs.min_val, obs.max_val) # check alpha value self.assertEqual(alpha, alpha_test) # calculate expected gamma value gamma_test = 0 for i in range(2): gamma_test += 2**(-i) gamma_test = 1 / gamma_test # check gamma value self.assertEqual(gamma, gamma_test) # check quantization levels size quantlevels_size_test = int(len(quantization_levels)) self.assertEqual(quantlevels_size_test, 49) # check negatives of each element contained # in quantization levels quantlevels_test_list = quantization_levels.tolist() negatives_contained = True for ele in quantlevels_test_list: if -ele not in quantlevels_test_list: negatives_contained = False self.assertTrue(negatives_contained) # check level indices size levelindices_size_test = int(len(level_indices)) self.assertEqual(levelindices_size_test, 49) # check level indices unique elements level_indices_test_list = level_indices.tolist() self.assertEqual(len(level_indices_test_list), len(set(level_indices_test_list))) """ Test case 5: calculate_qparams Assume hardcoded parameters: * b = 6 (total number of bits across all terms) * k = 1 (base bitwidth, i.e. bitwidth of every term) * n = 6 (number of additive terms) """ def test_calculate_qparams_k1(self): obs = APoTObserver(b=6, k=1) obs.min_val = torch.tensor([0.0]) obs.max_val = torch.tensor([1.0]) alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=False) # calculate expected gamma value gamma_test = 0 for i in range(6): gamma_test += 2**(-i) gamma_test = 1 / gamma_test # check gamma value self.assertEqual(gamma, gamma_test) # check quantization levels size quantlevels_size_test = int(len(quantization_levels)) quantlevels_size = 2**6 self.assertEqual(quantlevels_size_test, quantlevels_size) # check level indices size levelindices_size_test = int(len(level_indices)) level_indices_size = 2**6 self.assertEqual(levelindices_size_test, level_indices_size) # check level indices unique values level_indices_test_list = level_indices.tolist() self.assertEqual(len(level_indices_test_list), len(set(level_indices_test_list))) """ Test forward method on hard-coded tensor with arbitrary values. Checks that alpha is max of abs value of max and min values in tensor. """ def test_forward(self): obs = APoTObserver(b=4, k=2) X = torch.tensor([0.0, -100.23, -37.18, 3.42, 8.93, 9.21, 87.92]) X = obs.forward(X) alpha, gamma, quantization_levels, level_indices = obs.calculate_qparams(signed=True) min_val = torch.min(X) max_val = torch.max(X) expected_alpha = torch.max(-min_val, max_val) self.assertEqual(alpha, expected_alpha) if __name__ == '__main__': unittest.main()