def test_max(self):
torch.manual_seed(12345)
ref_lenet = QuantLeNet()
torch.manual_seed(12345)
test_lenet = QuantLeNet()
for module in ref_lenet.modules():
if isinstance(module,
(quant_nn.QuantConv2d, quant_nn.QuantLinear)):
module.weight_quantizer.enable_calib()
module.weight_quantizer.disable_quant()
module.weight_quantizer(module.weight)
module.weight_quantizer.load_calib_amax()
calib.calibrate_weights(test_lenet, method="max")
for ref_module, test_module in zip(ref_lenet.modules(),
test_lenet.modules()):
if isinstance(ref_module,
(quant_nn.QuantConv2d, quant_nn.QuantLinear)):
test_utils.compare(ref_module.weight_quantizer.amax,
test_module.weight_quantizer.amax,
rtol=0,
atol=0,
ctol=0)
assert ref_module.weight_quantizer.amax.shape == test_module.weight_quantizer.amax.shape
def test_data_augmentation_collapse():
dataset = DatasetImpl(path=data_path,
shape=[256, 256],
augmentation=True,
collapse_length=2,
is_raw=True)
gt_idx, gt_k, gt_flip, gt_angle = 1, 2, False, 0
gt_box, gt_seq_length = np.array([0, 0, 260, 346]), 1
events, timestamps, images, aug_params = dataset.__getitem__(
idx=gt_idx,
k=gt_k,
is_flip=gt_flip,
angle=gt_angle,
box=gt_box,
seq_length=gt_seq_length)
assert gt_idx == aug_params[0]
assert gt_seq_length == aug_params[1]
assert gt_k == aug_params[2]
assert (gt_box == aug_params[3]).all()
assert gt_angle == aug_params[4]
assert gt_flip == aug_params[5]
element1 = tuple(read_test_elem(1, element_index=0, box=gt_box))
element2 = tuple(read_test_elem(2, element_index=0, box=gt_box))
gt_events = concat_events(element1[0], element2[0])
gt_timestamps = np.array([0, element2[2] - element1[1]])
gt_events['timestamp'] -= element1[1]
assert element1[2] == element2[1]
assert (element1[4] == element2[3]).all()
gt_images = np.concatenate([element1[3][None], element2[4][None]],
axis=0).astype(np.float32)
compare(events, gt_events)
assert (timestamps == gt_timestamps).all()
assert (images == gt_images).all()
def test_cuda_ext(self):
x_np = np.random.rand(1023).astype('float32')
x_torch = torch.Tensor(x_np).cuda()
for num_bits in [3, 4, 5, 7, 8, 11]:
for unsigned in [True, False]:
test_utils.compare(cuda_ext.fake_tensor_quant(
x_torch, torch.max(torch.abs(x_torch)), num_bits,
unsigned),
tensor_quant.fake_tensor_quant(
x_torch, torch.max(torch.abs(x_torch)),
num_bits, unsigned),
rtol=0,
atol=0)
# Test fp16
x_np_fp16 = np.random.rand(1023).astype('float16')
x_torch_fp16 = torch.Tensor(x_np_fp16).cuda().half()
test_utils.compare(
cuda_ext.fake_tensor_quant(x_torch_fp16,
torch.max(torch.abs(x_torch_fp16))),
tensor_quant.fake_tensor_quant(x_torch_fp16,
torch.max(torch.abs(x_torch_fp16))),
rtol=0,
atol=0)
def test_against_unquantized(self):
kernel_size = 3
test_input = torch.randn(16, _NUM_IN_CHANNELS, 24, 24, 24).cuda()
torch.manual_seed(1234)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(1234)
fake_quant_conv3d = quant_conv.QuantConv3d(
_NUM_IN_CHANNELS,
_NUM_OUT_CHANNELS,
kernel_size,
bias=True,
quant_desc_input=QuantDescriptor(num_bits=16),
quant_desc_weight=QuantDescriptor(num_bits=16, axis=(0)))
# Reset seed. Make sure weight and bias are the same
torch.manual_seed(1234)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(1234)
conv3d = nn.Conv3d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
fake_quant_output = fake_quant_conv3d(test_input)
output = conv3d(test_input)
test_utils.compare(fake_quant_output, output, rtol=1e-6, atol=2e-4)
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_track_amax(self):
max_calibrator = calib.MaxCalibrator(8, None, False, track_amax=True)
x_1 = torch.rand(129).cuda()
x_2 = torch.rand(127).cuda()
max_calibrator.collect(x_1)
max_calibrator.collect(x_2)
test_utils.compare(max_calibrator.compute_amax(), torch.max(x_1.max(), x_2.max()), atol=0, rtol=0, ctol=0)
np.testing.assert_array_equal(max_calibrator.amaxs[0], x_1.max().cpu().numpy())
np.testing.assert_array_equal(max_calibrator.amaxs[1], x_2.max().cpu().numpy())
def test_simple_run(self):
max_calibrator = calib.MaxCalibrator(8, None, False)
x_1 = torch.rand(129).cuda()
x_2 = torch.rand(127).cuda()
max_calibrator.collect(x_1)
max_calibrator.collect(x_2)
test_utils.compare(max_calibrator.compute_amax(), torch.max(x_1.max(), x_2.max()), atol=0, rtol=0, ctol=0)
# Nothing to test other than creation
max_calibrator = calib.MaxCalibrator(8, None, True)
def test_cuda_ext_with_axis(self):
x_np = np.random.rand(3, 4, 5, 6).astype('float32')
x_torch = torch.Tensor(x_np).cuda()
# amax along axis 1
amax_torch = torch.tensor([0.8, 0.9, 0.7, 0.6], device="cuda")
for num_bits in [3, 4, 5, 7, 8, 11]:
for unsigned in [True, False]:
cuda_ext_out = cuda_ext.fake_tensor_quant_with_axis(
x_torch, amax_torch, 1, num_bits, unsigned)
pytorch_out = tensor_quant.fake_tensor_quant(
x_torch, amax_torch.view(1, -1, 1, 1), num_bits, unsigned)
test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)
def test_fine_grain(self):
axis = 0
reducs_axis = (1, 2, 3)
max_calibrator = calib.MaxCalibrator(8, axis, False)
x_1 = torch.rand(31, 63, 7, 7).cuda()
x_2 = torch.rand(31, 63, 7, 7).cuda()
max_calibrator.collect(x_1)
max_calibrator.collect(x_2)
assert max_calibrator.compute_amax().shape[0] == 31
test_utils.compare(max_calibrator.compute_amax(),
quant_utils.reduce_amax(torch.max(x_1, x_2), axis=reducs_axis),
atol=0, rtol=0, ctol=0)
max_calibrator.reset()
assert max_calibrator.compute_amax() is None
def test_entropy_and_percentile_calib(self):
"""Don't really have a good way to test it."""
quant_desc1 = tensor_quant.QuantDescriptor(calib_method='histogram')
quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1, if_calib=True, if_quant=False).cuda()
x_1 = torch.rand(3, 63, 7, 7).cuda()
x_2 = torch.rand(3, 63, 7, 7).cuda()
quantizer1(x_1)
quantizer1(x_2)
quantizer1.load_calib_amax("entropy")
test_utils.compare(quantizer1._calibrator.compute_amax("entropy"), quantizer1.amax, atol=0, rtol=0, ctol=0)
quantizer1._calibrator.reset()
quantizer1(x_1)
quantizer1(x_2)
quantizer1.load_calib_amax("percentile", percentile=99.99)
test_utils.compare(quantizer1._calibrator.compute_amax(
"percentile", percentile=99.99), quantizer1.amax, atol=0, rtol=0, ctol=0)
def test_dataloader():
dataset = DatasetImpl(path=data_path,
shape=[260, 346],
augmentation=False,
collapse_length=1,
is_raw=True)
data_loader = torch.utils.data.DataLoader(dataset,
collate_fn=collate_wrapper,
batch_size=2,
pin_memory=True,
shuffle=False)
batch = next(iter(data_loader))
element1 = tuple(read_test_elem(0, element_index=0, is_torch=True))
element2 = tuple(read_test_elem(1, element_index=0, is_torch=True))
element1[0]['timestamp'] -= element1[1]
element2[0]['timestamp'] -= element2[1]
gt_events = concat_events(element1[0], element2[0])
gt_events['sample_index'] = np.hstack(
[np.full_like(element1[0]['x'], 0),
np.full_like(element2[0]['x'], 1)])
gt_events = {k: torch.tensor(v) for k, v in gt_events.items()}
gt_timestamps = torch.tensor(
[0, element1[2] - element1[1], 0, element2[2] - element2[1]],
dtype=torch.float32)
gt_sample_idx = torch.tensor([0, 0, 1, 1], dtype=torch.long)
image00 = torch.tensor(element1[3], dtype=torch.float32)[None, None]
image01 = torch.tensor(element1[4], dtype=torch.float32)[None, None]
image10 = torch.tensor(element2[3], dtype=torch.float32)[None, None]
image11 = torch.tensor(element2[4], dtype=torch.float32)[None, None]
gt_images = torch.cat([image00, image01, image10, image11], dim=0) \
.to(torch.float32)
compare(batch['events'], gt_events)
assert torch.equal(batch['timestamps'], gt_timestamps)
assert torch.equal(batch['sample_idx'], gt_sample_idx)
assert (batch['images'] == gt_images).all()
assert batch['size'] == 2
def test_max_calib(self):
axis = 0
reduce_axis = (1, 2, 3)
quant_desc1 = tensor_quant.QuantDescriptor(axis=axis)
quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1).cuda()
quantizer1.enable_calib()
quant_desc1 = tensor_quant.QuantDescriptor(axis=axis)
quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1).cuda()
quantizer1.enable_calib()
with pytest.raises(RuntimeError, match="Calibrator returned None"):
quantizer1.load_calib_amax()
x_1 = torch.rand(127, 63, 7, 7).cuda()
x_2 = torch.rand(127, 63, 7, 7).cuda()
quantizer1(x_1)
quantizer1(x_2)
quantizer1.disable_calib()
global_amax = torch.max(
quant_utils.reduce_amax(x_1, axis=reduce_axis, keepdims=True),
quant_utils.reduce_amax(x_2, axis=reduce_axis, keepdims=True))
test_utils.compare(quantizer1._calibrator.compute_amax(), global_amax, atol=0, rtol=0, ctol=0)
quantizer1.load_calib_amax()
test_utils.compare(quantizer1.amax, global_amax, atol=0, rtol=0, ctol=0)
quant_desc2 = tensor_quant.QuantDescriptor(learn_amax=True)
quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2).cuda()
quantizer2.enable_calib()
quantizer2(x_1)
quantizer2(x_2)
quantizer2.load_calib_amax()
quantizer2.init_learn_amax()
test_utils.compare(quantizer2.clip.clip_value_min, -torch.max(global_amax), atol=0, rtol=0, ctol=0)
test_utils.compare(quantizer2.clip.clip_value_max, torch.max(global_amax), atol=0, rtol=0, ctol=0)
def test_read_write(self):
with tempfile.NamedTemporaryFile(suffix='.hdf5') as f:
filename = Path(f.name)
write_encoded_batch(filename, self.encoded_batch)
assert filename.is_file()
with h5py.File(filename, 'r') as f:
channels_per_sample = torch.tensor(f['channels_per_sample'])
elements_per_sample = torch.tensor(f['elements_per_sample'])
read = read_encoded_quantized_batch(f, channels_per_sample,
elements_per_sample, 0, 3)
compare(read, self.encoded_batch)
with h5py.File(filename, 'r') as f:
channels_per_sample = torch.tensor(f['channels_per_sample'])
elements_per_sample = torch.tensor(f['elements_per_sample'])
read = read_encoded_quantized_batch(f, channels_per_sample,
elements_per_sample, 0, 2)
compare(read, self.encoded_batches[0])
with h5py.File(filename, 'r') as f:
channels_per_sample = torch.tensor(f['channels_per_sample'])
elements_per_sample = torch.tensor(f['elements_per_sample'])
read = read_encoded_quantized_batch(f, channels_per_sample,
elements_per_sample, 2, 3)
compare(read, self.encoded_batches[1])
filename.unlink()
def test_encode(self):
encoded = encode_quantized_batch(self.decoded_batch)
compare(encoded, self.encoded_batch)
def test_join(self):
joined = join_batches(self.encoded_batches)
compare(joined, self.encoded_batch)
def test_decode(self):
decoded = decode_quantized_batch(self.encoded_batch)
compare(decoded, self.decoded_batch)
