chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,278 @@
|
||||
#
|
||||
# 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)
|
||||
Reference in New Issue
Block a user