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)
