def test_grow(self, verbose):
x_1 = torch.tensor([0, 255, 255, 255, 255, 255]).cuda()
x_2 = torch.tensor([0, 255, 255, 255, 255, 256]).cuda()
hist_calibrator = calib.HistogramCalibrator(8,
None,
False,
grow_method='stretch')
hist_calibrator.collect(x_1)
hist_calibrator.collect(x_2)
amax = hist_calibrator.compute_amax(method='entropy')
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be closer to 256 because the last bin gets stretched to (~255, 257)
assert (amax - 255.).abs() < (amax - 256.).abs()
hist_calibrator = calib.HistogramCalibrator(8,
None,
False,
grow_method='append')
hist_calibrator.collect(x_1)
hist_calibrator.collect(x_2)
amax = hist_calibrator.compute_amax(method='mse')
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be closer to 255
assert (amax - 255.).abs() < 0.5
def test_torch_hist(self):
x_1 = torch.rand(1023, device="cuda")
x_1[0] = 0
x_2 = torch.rand(
1023,
device="cuda") + 1 # Make sure histogram bins need to be grown
x_2[1] = 0
calibrator_np = calib.HistogramCalibrator(8,
None,
False,
num_bins=19,
torch_hist=False)
calibrator_torch = calib.HistogramCalibrator(8,
None,
False,
num_bins=19,
torch_hist=True)
calibrator_np.collect(x_1)
calibrator_torch.collect(x_1)
assert calibrator_torch._calib_hist.numel(
) == calibrator_torch._calib_bin_edges.numel() - 1
np.testing.assert_array_equal(
calibrator_np._calib_hist,
calibrator_torch._calib_hist.cpu().numpy())
np.testing.assert_array_almost_equal(
calibrator_np._calib_bin_edges,
calibrator_torch._calib_bin_edges.cpu().numpy())
# Test multiple collections with some of them needs to expand range
for _ in range(3):
calibrator_np.collect(x_2)
calibrator_torch.collect(x_2)
calibrator_np.collect(x_1)
calibrator_torch.collect(x_1)
# Test compute_amax function doesn't convert _calib_hist and _calib_bin_edges unnecessarily
calibrator_np.compute_amax("percentile", percentile=99.99)
calibrator_torch.compute_amax("percentile", percentile=99.99)
np.testing.assert_array_equal(
calibrator_np._calib_hist,
calibrator_torch._calib_hist.cpu().numpy())
np.testing.assert_array_almost_equal(
calibrator_np._calib_bin_edges,
calibrator_torch._calib_bin_edges.cpu().numpy())
assert calibrator_torch._calib_hist.numel(
) == calibrator_torch._calib_bin_edges.numel() - 1
def test_range(self):
calibrator = calib.HistogramCalibrator(8, None, False)
x_1 = torch.arange(100)
calibrator.collect(x_1)
with pytest.raises(ValueError, match="range"):
calibrator.compute_amax("percentile", percentile=-10)
with pytest.raises(ValueError, match="range"):
calibrator.compute_amax("percentile", percentile=200)
def test_mse_with_axis(self):
torch.manual_seed(12345)
test_lenet = QuantLeNet()
ref_calibrator = calib.HistogramCalibrator(8, None, False)
calib.calibrate_weights(test_lenet, method="mse", perchannel=True)
ref_calibrator.collect(test_lenet.conv2.weight[1])
ref_amax = ref_calibrator.compute_amax("mse")
test_utils.compare(ref_amax, test_lenet.conv2.weight_quantizer.amax[1], rtol=0, atol=0, ctol=0)
def test_percentile(self):
torch.manual_seed(12345)
test_lenet = QuantLeNet()
test_percentile = 99.99
ref_calibrator = calib.HistogramCalibrator(8, None, False)
calib.calibrate_weights(test_lenet, method="percentile", perchannel=False, percentile=test_percentile)
ref_calibrator.collect(test_lenet.conv1.weight)
ref_amax = ref_calibrator.compute_amax("percentile", percentile=test_percentile)
test_utils.compare(ref_amax, test_lenet.conv1.weight_quantizer.amax, rtol=0, atol=0, ctol=0)
def test_unsigned_one_tensor(self, verbose):
calibrator = calib.HistogramCalibrator(8, None, True)
x_1 = torch.arange(100)
calibrator.collect(x_1)
amax = calibrator.compute_amax("percentile", percentile=80)
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be approximately 79
assert (amax - 79.).abs() < 100 / 2048
def test_unsigned(self, verbose):
hist_calibrator = calib.HistogramCalibrator(8, None, True, grow_method='stretch')
x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
x_2[1, 1, 1, 1] = 10. # create outlier
hist_calibrator.collect(x_2)
amax = hist_calibrator.compute_amax("entropy")
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
assert amax < 1.1
def test_unsigned_one_tensor(self, verbose):
calibrator = calib.HistogramCalibrator(8, None, True)
x_1 = torch.ones(11, 7, 3, 3).cuda() * 512.
x_1[1, 1, 1, 1] = 513. # create an outlier
calibrator.collect(x_1)
amax = calibrator.compute_amax("mse")
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be closer to 512
assert (amax - 512.).abs() < (amax - 513.).abs()
def test_two_tensor(self, verbose):
calibrator = calib.HistogramCalibrator(8, None, False)
x_1 = torch.arange(100)
calibrator.collect(x_1)
x_2 = torch.arange(0, 50, 0.5)
calibrator.collect(x_2)
amax = calibrator.compute_amax("percentile", percentile=99)
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be approximately 97
assert (amax - 97.).abs() < 100/1024
def test_one_tensor(self, verbose):
hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='stretch')
x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
x_2[1, 1, 1, 1] = 10. # create outlier
hist_calibrator.collect(x_2)
# Don't have a better test metric. One outlier 10 should be discared by KL-divergence
amax = hist_calibrator.compute_amax("entropy")
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
assert amax < 1.1
def __init__(self,
quant_desc=QuantDescriptor(),
disabled=False,
if_quant=True,
if_clip=False,
if_calib=False):
"""Initialize quantizer and set up required variables"""
super(TensorQuantizer, self).__init__()
# Expand quant_desc. Use quant_desc.dict would be eaiser, but adding one-by-one explicitly gives more control
self._num_bits = quant_desc.num_bits
self._fake_quant = quant_desc.fake_quant
self._axis = quant_desc.axis
self._scale_amax = quant_desc.scale_amax
self._learn_amax = quant_desc.learn_amax
self._unsigned = quant_desc.unsigned
self._narrow_range = quant_desc.narrow_range
self._scale = None if not quant_desc.fake_quant else 1.
self._disabled = disabled
self._if_quant = if_quant
self._if_clip = False
self._if_calib = if_calib
if quant_desc.amax is not None:
self.register_buffer('_amax', torch.tensor(quant_desc.amax))
# Clip module consumes a lot of memory, so only create it if learn_amax is True
if self._learn_amax:
init_amax = quant_desc.amax if quant_desc.amax is not None else 1.
self.clip = Clip(-init_amax,
init_amax,
learn_min=True,
learn_max=True)
# It makes more sense to enable clip stage (which learns amax) if learn_amax is true
self.enable_clip()
if if_clip:
self.enable_clip()
if quant_desc.calib_method == "histogram":
logging.info("Creating histogram calibrator")
self._calibrator = calib.HistogramCalibrator(
num_bits=self._num_bits,
axis=self._axis,
unsigned=self._unsigned)
elif quant_desc.calib_method == "max":
logging.info("Creating Max calibrator")
self._calibrator = calib.MaxCalibrator(num_bits=self._num_bits,
axis=self._axis,
unsigned=self._unsigned)
def test_skip_zeros(self, verbose):
x_1 = torch.tensor([0, 0, 0, 0, 0, 1, 2, 3, 4, 5])
x_2 = torch.tensor([0, 0, 0, 0, 0, 6, 7, 8, 9, 10])
calibrator = calib.HistogramCalibrator(8, None, False, skip_zeros=True)
calibrator.collect(x_1)
calibrator.collect(x_2)
amax = calibrator.compute_amax("percentile", percentile=50)
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be close to 5
assert (amax - 5.).abs() < 10 / 2048
def test_two_tensor(self, verbose):
calibrator = calib.HistogramCalibrator(8, None, False)
x_1 = torch.ones(11, 7, 3, 3).cuda() * 255.
x_1[1, 1, 1, 1] = 256. # create an outlier
calibrator.collect(x_1)
x_2 = torch.ones(11, 7, 3, 3).cuda() * 255.
calibrator.collect(x_2)
amax = calibrator.compute_amax("mse")
if verbose:
print('amax={:.4f}'.format(amax.item()), end=' ')
# amax should be closer to 255
assert (amax - 255.).abs() < (amax - 256.).abs()
def test_print_calibrator(self):
print(calib.MaxCalibrator(7, 1, False))
hist_calibrator = calib.HistogramCalibrator(8, None, True)
hist_calibrator.collect(torch.rand(10))
print(hist_calibrator)
def test_repr(self):
calibrator = calib.HistogramCalibrator(8, None, False)
repr(calibrator)
def test_repr(self):
hist_calibrator = calib.HistogramCalibrator(8, None, True)
repr(hist_calibrator)
