def test_learnable_backward_per_channel_cuda(self, X):
torch.random.manual_seed(NP_RANDOM_SEED)
X, (scale, zero_point, axis, torch_type) = X
X_base = torch.tensor(X).to('cuda')
scale_base = to_tensor(scale, 'cuda')
zero_point_base = to_tensor(zero_point, 'cuda')
self._test_learnable_backward_per_channel(X_base, 'cuda', scale_base,
zero_point_base, axis)
def test_forward_per_channel(self, device, X):
r"""Tests the forward path of the FakeQuantizePerTensorAffine op.
"""
np.random.seed(NP_RANDOM_SEED)
X, (scale, zero_point, axis, torch_type) = X
quant_min = torch.iinfo(torch_type).min
quant_max = torch.iinfo(torch_type).max
X = to_tensor(X, device)
scale = to_tensor(scale, device)
zero_point = torch.tensor(zero_point).to(dtype=torch.int32, device=device)
Y = _fake_quantize_per_channel_affine_reference(X.cpu(), scale.cpu(), zero_point.cpu(), axis, quant_min, quant_max)
Y_prime = torch.fake_quantize_per_channel_affine(
X, scale, zero_point, axis, quant_min, quant_max)
np.testing.assert_allclose(Y, Y_prime.cpu(), rtol=tolerance, atol=tolerance)
def test_fq_module_per_channel(self, device, X):
np.random.seed(NP_RANDOM_SEED)
X, (scale, zero_point, axis, torch_type) = X
quant_min = torch.iinfo(torch_type).min
quant_max = torch.iinfo(torch_type).max
X = to_tensor(X, device)
X.requires_grad_()
fq_module = FakeQuantize(default_per_channel_weight_observer,
quant_min,
quant_max,
ch_axis=axis).to(device)
Y_prime = fq_module(X)
assert fq_module.scale is not None
assert fq_module.zero_point is not None
Y = _fake_quantize_per_channel_affine_reference(
X, fq_module.scale, fq_module.zero_point, axis, quant_min,
quant_max)
np.testing.assert_allclose(Y.cpu().detach().numpy(),
Y_prime.cpu().detach().numpy(),
rtol=tolerance,
atol=tolerance)
# Test backward
dout = torch.rand_like(X, dtype=torch.float, device=device)
Y_prime.backward(dout)
dX = _fake_quantize_per_channel_affine_grad_reference(
dout, X, fq_module.scale, fq_module.zero_point, axis, quant_min,
quant_max)
np.testing.assert_allclose(dX.cpu().numpy(),
X.grad.cpu().detach().numpy(),
rtol=tolerance,
atol=tolerance)
def test_backward_per_channel(self, device, X):
r"""Tests the backward method.
"""
np.random.seed(NP_RANDOM_SEED)
X, (scale, zero_point, axis, torch_type) = X
quant_min = torch.iinfo(torch_type).min
quant_max = torch.iinfo(torch_type).max
X = to_tensor(X, device)
scale = to_tensor(scale, device)
zero_point = torch.tensor(zero_point).to(dtype=torch.int32, device=device)
X.requires_grad_()
Y_prime = torch.fake_quantize_per_channel_affine(
X, scale, zero_point, axis, quant_min, quant_max)
dout = torch.rand_like(X, dtype=torch.float).to(device)
dX = _fake_quantize_per_channel_affine_grad_reference(
dout, X, scale, zero_point, axis, quant_min, quant_max)
Y_prime.backward(dout)
np.testing.assert_allclose(dX.cpu().detach().numpy(), X.grad.cpu().detach().numpy(), rtol=tolerance, atol=tolerance)
def test_fixed_qparams_fq_module(self, device, X):
X, (scale, zero_point, torch_type) = X
X = to_tensor(X, device)
fq_module = default_affine_fixed_qparams_fake_quant()
fixed_scale = fq_module.scale.clone()
fixed_zero_point = fq_module.zero_point.clone()
# run fq module and make sure the quantization parameters does not change
torch.quantization.enable_observer(fq_module)
fq_module(X)
self.assertEqual(fixed_scale, fq_module.scale)
self.assertEqual(fixed_zero_point, fq_module.zero_point)
def test_numerical_consistency_per_channel(self, device, X):
r"""Comparing numerical consistency between CPU quantize/dequantize op and the CPU fake quantize op
"""
np.random.seed(NP_RANDOM_SEED)
X, (scale, zero_point, axis, torch_type) = X
quant_min = torch.iinfo(torch_type).min
quant_max = torch.iinfo(torch_type).max
X = to_tensor(X, device)
scale = to_tensor(scale, device)
zero_point = torch.tensor(zero_point).to(dtype=torch.int64,
device=device)
# quantize_linear and dequantize are only implemented in CPU
Y = torch.dequantize(
torch.quantize_per_channel(X.cpu(), scale.cpu(), zero_point.cpu(),
axis, torch_type))
Y_prime = torch.fake_quantize_per_channel_affine(
X, scale, zero_point, axis, quant_min, quant_max)
np.testing.assert_allclose(Y,
Y_prime.cpu(),
rtol=tolerance,
atol=tolerance)
