def test_torch_quantizer_modules_detection(self): # test if modules can be detected model = TorchModel() config_list = [{ 'quant_types': ['weight'], 'quant_bits': 8, 'op_types': ['Conv2d', 'Linear'] }, { 'quant_types': ['output'], 'quant_bits': 8, 'quant_start_step': 0, 'op_types': ['ReLU'] }] model.relu = torch.nn.ReLU() quantizer = torch_compressor.QAT_Quantizer(model, config_list) quantizer.compress() modules_to_compress = quantizer.get_modules_to_compress() modules_to_compress_name = [t[0].name for t in modules_to_compress] assert "conv1" in modules_to_compress_name assert "conv2" in modules_to_compress_name assert "fc1" in modules_to_compress_name assert "fc2" in modules_to_compress_name assert "relu" in modules_to_compress_name assert len(modules_to_compress_name) == 5
def test_torch_QAT_quantizer(self): model = TorchModel() config_list = [{ 'quant_types': ['weight'], 'quant_bits': 8, 'op_types': ['Conv2d', 'Linear'] }, { 'quant_types': ['output'], 'quant_bits': 8, 'quant_start_step': 0, 'op_types': ['ReLU'] }] model.relu = torch.nn.ReLU() quantizer = torch_compressor.QAT_Quantizer(model, config_list) quantizer.compress() # test quantize # range not including 0 eps = 1e-7 weight = torch.tensor([[1, 2], [3, 5]]).float() quantize_weight = quantizer.quantize_weight(weight, config_list[0], model.conv2) assert math.isclose(model.conv2.scale, 5 / 255, abs_tol=eps) assert model.conv2.zero_point == 0 # range including 0 weight = torch.tensor([[-1, 2], [3, 5]]).float() quantize_weight = quantizer.quantize_weight(weight, config_list[0], model.conv2) assert math.isclose(model.conv2.scale, 6 / 255, abs_tol=eps) assert model.conv2.zero_point in (42, 43) # test ema x = torch.tensor([[-0.2, 0], [0.1, 0.2]]) out = model.relu(x) assert math.isclose(model.relu.module.tracked_min_biased, 0, abs_tol=eps) assert math.isclose(model.relu.module.tracked_max_biased, 0.002, abs_tol=eps) quantizer.step() x = torch.tensor([[0.2, 0.4], [0.6, 0.8]]) out = model.relu(x) assert math.isclose(model.relu.module.tracked_min_biased, 0.002, abs_tol=eps) assert math.isclose(model.relu.module.tracked_max_biased, 0.00998, abs_tol=eps)
def test_torch_QAT_quantizer(self): model = TorchModel() config_list = [{ 'quant_types': ['weight'], 'quant_bits': 8, 'op_types': ['Conv2d', 'Linear'] }, { 'quant_types': ['output'], 'quant_bits': 8, 'quant_start_step': 0, 'op_types': ['ReLU'] }] model.relu = torch.nn.ReLU() quantizer = torch_compressor.QAT_Quantizer(model, config_list) quantizer.compress() # test quantize # range not including 0 eps = 1e-7 weight = torch.tensor([[1, 2], [3, 5]]).float() model.conv2.module.old_weight.data = weight quantizer.quantize_weight(model.conv2) assert math.isclose(model.conv2.module.scale, 5 / 255, abs_tol=eps) assert model.conv2.module.zero_point == 0 # range including 0 weight = torch.tensor([[-1, 2], [3, 5]]).float() model.conv2.module.old_weight.data = weight quantizer.quantize_weight(model.conv2) assert math.isclose(model.conv2.module.scale, 6 / 255, abs_tol=eps) assert model.conv2.module.zero_point in (42, 43) # test value of weight and bias after quantization weight = torch.tensor([[1.1287, 2.3456], [3.7814, 5.9723]]) weight_valid = torch.tensor([[1.1242, 2.3421], [3.7707, 5.9723]]) bias = torch.tensor([2.3432, 3.4342, 1.3414, 5.2341]) bias_valid = torch.tensor([2.3432, 3.4342, 1.3414, 5.2341]) model.conv2.module.old_weight.data = weight model.conv2.module.bias.data = bias quantizer.quantize_weight(model.conv2) assert torch.all( torch.isclose(model.conv2.module.weight.data, weight_valid, rtol=1e-4)) assert torch.all( torch.isclose(model.conv2.module.bias.data, bias_valid, rtol=1e-7)) # test ema eps = 1e-7 x = torch.tensor([[-0.2, 0], [0.1, 0.2]]) out = model.relu(x) assert math.isclose(model.relu.module.tracked_min_biased, 0, abs_tol=eps) assert math.isclose(model.relu.module.tracked_max_biased, 0.002, abs_tol=eps) quantizer.step_with_optimizer() x = torch.tensor([[0.2, 0.4], [0.6, 0.8]]) out = model.relu(x) assert math.isclose(model.relu.module.tracked_min_biased, 0.002, abs_tol=eps) assert math.isclose(model.relu.module.tracked_max_biased, 0.00998, abs_tol=eps)