wehub-resource-sync
278 lines
12 KiB
Python
278 lines
12 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
"""tests of tensor quantizer"""
|
|
import contextlib
|
|
|
|
import pytest
|
|
import numpy as np
|
|
|
|
import torch
|
|
|
|
from pytorch_quantization import tensor_quant
|
|
from pytorch_quantization.nn.modules import tensor_quantizer
|
|
from pytorch_quantization import utils as quant_utils
|
|
import tests.utils as test_utils
|
|
from tests.fixtures import verbose
|
|
|
|
np.random.seed(12345)
|
|
|
|
# pylint:disable=missing-docstring, no-self-use
|
|
|
|
|
|
class TestTensorQuantizer():
|
|
|
|
def test_simple_run(self):
|
|
"""Quantizer calls fake_tensor_quant by default"""
|
|
x = torch.randn(3, 7).cuda()
|
|
amax_x = torch.max(torch.abs(x))
|
|
fn_quant_x = tensor_quant.fake_tensor_quant(x, amax_x)
|
|
quantizer = tensor_quantizer.TensorQuantizer()
|
|
module_quant_x = quantizer(x)
|
|
np.testing.assert_array_equal(fn_quant_x.cpu().numpy(), module_quant_x.cpu().numpy())
|
|
|
|
def test_simple_run_no_fake(self):
|
|
"""Quantizer fake_quant=False calls tensor_quant and sets the scale property"""
|
|
x = torch.randn(3, 7).cuda()
|
|
amax_x = torch.max(torch.abs(x))
|
|
fn_quant_x, fn_scale = tensor_quant.tensor_quant(x, amax_x)
|
|
quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(num_bits=8, fake_quant=False))
|
|
module_quant_x = quantizer(x)
|
|
module_scale = quantizer.scale
|
|
np.testing.assert_array_equal(fn_quant_x.cpu().numpy(), module_quant_x.cpu().numpy())
|
|
np.testing.assert_array_equal(fn_scale.cpu().numpy(), module_scale.cpu().numpy())
|
|
|
|
def test_per_tensor_scale(self):
|
|
"""Quantizer performs expected quantization"""
|
|
x_np = np.random.rand(1023)
|
|
x_torch = torch.Tensor(x_np)
|
|
quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)))
|
|
quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(num_bits=8, fake_quant=False))
|
|
module_quant_x = quantizer(x_torch)
|
|
np.testing.assert_array_equal(module_quant_x.cpu().numpy(), quant_x_np)
|
|
|
|
def test_per_channel_scale(self, verbose):
|
|
"""Quantizer performs per channel scaling"""
|
|
x_np = np.random.rand(15, 15, 64, 128).astype('float32')
|
|
x_torch = torch.Tensor(x_np).cuda()
|
|
|
|
# Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
|
|
# Shrink max a little, so that clip behavior is tested
|
|
amax_x_np = 0.7 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
|
|
|
|
quant_x_np = test_utils.quant_np(x_np, amax_x_np)
|
|
quantizer = tensor_quantizer.TensorQuantizer(
|
|
tensor_quant.QuantDescriptor(num_bits=8, axis=(0), fake_quant=False, scale_amax=0.7))
|
|
quantizer.cuda()
|
|
module_quant_x = quantizer(x_torch)
|
|
|
|
# np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
|
|
# Pytorch numerics is not the same as numpy, it will be off by 1
|
|
error = np.abs(module_quant_x.cpu().numpy() - quant_x_np)
|
|
np.testing.assert_array_less(error, 2)
|
|
if verbose:
|
|
mismatches = np.where(error >= 1)
|
|
print("Mismatches:")
|
|
print(" Original: ", x_np[mismatches])
|
|
print(" numpy: ", quant_x_np[mismatches])
|
|
print(" TensorQuantizer: ", module_quant_x.cpu().numpy()[mismatches])
|
|
|
|
def test_learn_amax(self):
|
|
"""Test the clip implied by learn_amax"""
|
|
x_np = np.random.rand(1023).astype(np.float32)
|
|
x_torch = torch.Tensor(x_np).cuda()
|
|
amax = 0.5
|
|
quant_x_np = test_utils.quant_np(x_np, 0.5, fake=True)
|
|
quantizer = tensor_quantizer.TensorQuantizer(
|
|
tensor_quant.QuantDescriptor(num_bits=8, amax=amax, learn_amax=True)).cuda()
|
|
assert hasattr(quantizer, 'clip')
|
|
module_quant_x = quantizer(x_torch)
|
|
np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
|
|
|
|
def test_clip_mode(self):
|
|
"""Test the clip stage only"""
|
|
x_np = np.random.rand(1023).astype(np.float32)
|
|
x_torch = torch.Tensor(x_np).cuda()
|
|
amax = 0.5
|
|
clip_x_np = np.clip(x_np, -amax, amax)
|
|
quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(amax=amax, learn_amax=True),
|
|
if_quant=False,
|
|
if_clip=True).cuda()
|
|
assert hasattr(quantizer, 'clip')
|
|
module_clip_x = quantizer(x_torch)
|
|
np.testing.assert_array_equal(module_clip_x.cpu().detach().numpy(), clip_x_np)
|
|
|
|
def test_scale_amax(self):
|
|
x_np = np.random.rand(1023).astype(np.float32)
|
|
x_torch = torch.Tensor(x_np).cuda()
|
|
amax = 0.5
|
|
scale_amax = 0.9
|
|
quant_x_np = test_utils.quant_np(x_np, amax * scale_amax, fake=True)
|
|
quantizer = tensor_quantizer.TensorQuantizer(
|
|
tensor_quant.QuantDescriptor(num_bits=8, amax=amax, scale_amax=scale_amax)).cuda()
|
|
module_quant_x = quantizer(x_torch)
|
|
np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
|
|
|
|
# Test twice. There was a but in scale amax logic that modify the amax every time
|
|
module_quant_x = quantizer(x_torch)
|
|
np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
|
|
|
|
def test_disable(self):
|
|
x = torch.randn(3, 7).cuda()
|
|
amax_x = torch.max(torch.abs(x))
|
|
quantizer = tensor_quantizer.TensorQuantizer(disabled=True).cuda()
|
|
module_quant_x = quantizer(x)
|
|
np.testing.assert_array_equal(x.cpu().numpy(), module_quant_x.cpu().numpy())
|
|
|
|
def test_state_loading(self):
|
|
"""Test quant_desc loading via state_dict"""
|
|
amax = [3.142, 2.718]
|
|
quant_desc1 = tensor_quant.QuantDescriptor(amax=amax)
|
|
quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1)
|
|
|
|
# copy state
|
|
quantizer1.load_state_dict(quantizer1.state_dict())
|
|
np.testing.assert_array_equal(quantizer1.amax.detach().cpu().numpy(), quant_desc1.amax)
|
|
|
|
def test_properties(self):
|
|
quant_desc1 = tensor_quant.QuantDescriptor(amax=3.14)
|
|
quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1)
|
|
quantizer1.amax = 0.577
|
|
|
|
assert quantizer1.amax.detach().cpu().numpy() == np.float32(0.577)
|
|
np.testing.assert_array_equal(quantizer1.amax.detach().cpu().numpy(), quantizer1.amax)
|
|
assert quantizer1.step_size == 0.577 / 127.
|
|
|
|
quant_desc2 = tensor_quant.QuantDescriptor()
|
|
quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2)
|
|
amax_np = np.array([3.142, 2.718], dtype=np.float32)
|
|
quantizer2.amax = amax_np
|
|
np.testing.assert_array_equal(quantizer2.amax.detach().cpu().numpy(), amax_np)
|
|
|
|
quant_desc3 = tensor_quant.QuantDescriptor()
|
|
quantizer3 = tensor_quantizer.TensorQuantizer(quant_desc3)
|
|
assert quantizer3.amax is None
|
|
|
|
def test_init_calib(self):
|
|
quant_desc2 = tensor_quant.QuantDescriptor(axis=(0, 1))
|
|
quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2, if_calib=True, if_quant=False).cuda()
|
|
|
|
x_2 = torch.rand(127, 63, 7, 7).cuda()
|
|
quantizer2(x_2)
|
|
quantizer2.load_calib_amax()
|
|
|
|
assert quantizer2.amax.numel() == 127 * 63
|
|
|
|
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_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_setters(self):
|
|
quantizer = tensor_quantizer.TensorQuantizer()
|
|
quantizer.num_bits = 7
|
|
quantizer.unsigned = True
|
|
|
|
assert quantizer.num_bits == 7
|
|
assert quantizer.unsigned
|
|
|
|
def test_pre_quant_scale(self):
|
|
quant_desc = tensor_quant.QuantDescriptor(axis=1, num_bits=8, amax=127.0)
|
|
quantizer = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
|
|
quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
|
|
|
|
inputs = torch.Tensor([[0, 0.4, 1.1, 2.0]]).cuda()
|
|
outputs_gt = torch.Tensor([[0, 0, 1, 2]]).cuda()
|
|
assert torch.allclose(quantizer(inputs), outputs_gt)
|
|
|
|
quantizer.pre_quant_scale = 2.0
|
|
outputs_gt = torch.Tensor([[0, 1, 2, 4]]).cuda()
|
|
assert torch.allclose(quantizer(inputs), outputs_gt)
|
|
|
|
quantizer2.pre_quant_scale = torch.Tensor([[1.0, 2.0, 3.0, 4.0]]).cuda()
|
|
outputs_gt = torch.Tensor([[0, 1, 3, 8]]).cuda()
|
|
assert torch.allclose(quantizer2(inputs), outputs_gt)
|
|
|
|
@pytest.mark.parametrize("E, M, axis", [(5, 2, None), (4, 3, None), (4, 3, 1), (7, 3, None)])
|
|
def test_e4m3(self, E, M, axis):
|
|
is_error_expected = (E != 4 or M != 3)
|
|
with (pytest.raises(TypeError)
|
|
if is_error_expected else contextlib.nullcontext()):
|
|
e4m3_desc = tensor_quant.QuantDescriptor(num_bits=(E, M), axis=axis)
|
|
e4m3_quantizer = tensor_quantizer.TensorQuantizer(e4m3_desc).to("cuda")
|
|
|
|
x = torch.rand(3, 63, 7, 7, device="cuda")
|
|
|
|
e4m3_x = e4m3_quantizer(x)
|
|
ref = tensor_quant.scaled_e4m3(x, e4m3_quantizer._get_amax(x), E, M)
|
|
test_utils.compare(e4m3_x, ref, atol=0, rtol=0) |