wehub-resource-sync
130 lines
4.9 KiB
Python
130 lines
4.9 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.
|
|
#
|
|
|
|
|
|
"""Utils for testing quantization."""
|
|
import numpy as np
|
|
from scipy.spatial import distance
|
|
|
|
import torch
|
|
|
|
from pytorch_quantization import tensor_quant
|
|
|
|
def quantize_by_range(x, num_bits):
|
|
"""Quantize torch tensor by range to num_bits with symmetric zero-mean quantizer."""
|
|
amax = x.abs().max()
|
|
x_q = tensor_quant.fake_tensor_quant(x, amax, num_bits)
|
|
return x_q
|
|
|
|
def quantize_by_range_fused(x_tuple, num_bits):
|
|
"""Quantize multiple torch tensors by combined range to num_bits with symmetric zero-mean quantizer."""
|
|
# compute aggregate amax across all tensors
|
|
amax = max([x.abs().max() for x in x_tuple])
|
|
# quantize each tensor with the aggregate amax
|
|
x_q_tuple = tuple(tensor_quant.fake_tensor_quant(x, amax, num_bits) for x in x_tuple)
|
|
return x_q_tuple
|
|
|
|
def copy_state_and_quantize(dst, src, num_bits):
|
|
"""Copy src to dst, quantize all 'weight' entries to num_bits."""
|
|
src_state_dict = src.state_dict()
|
|
dst_state_dict = dict()
|
|
for key in src_state_dict:
|
|
if 'weight' in key:
|
|
dst_state_dict[key] = quantize_by_range(src_state_dict[key], num_bits)
|
|
else:
|
|
dst_state_dict[key] = src_state_dict[key].clone()
|
|
|
|
dst.load_state_dict(dst_state_dict)
|
|
|
|
def copy_state_and_quantize_fused(dst, src, num_bits):
|
|
"""Copy src to dst, quantize all 'weight' entries to num_bits using the aggregate amax."""
|
|
src_state_dict = src.state_dict()
|
|
dst_state_dict = dict()
|
|
|
|
# compute aggregate amax across all weight tensors
|
|
amax = 0
|
|
for key in src_state_dict:
|
|
if 'weight' in key:
|
|
amax = max(amax, src_state_dict[key].abs().max())
|
|
|
|
# quantize each weight tensor with the aggregate amax
|
|
for key in src_state_dict:
|
|
if 'weight' in key:
|
|
dst_state_dict[key] = tensor_quant.fake_tensor_quant(src_state_dict[key], amax, num_bits)
|
|
else:
|
|
dst_state_dict[key] = src_state_dict[key].clone()
|
|
|
|
dst.load_state_dict(dst_state_dict)
|
|
|
|
def compare(a, b, rtol=1e-7, atol=1e-6, ctol=1e-6):
|
|
"""Compare two tensors and raise AssertionError if their difference is outside of tolerance."""
|
|
if torch.isinf(a).any():
|
|
raise ValueError("a contains infs")
|
|
if torch.isinf(b).any():
|
|
raise ValueError("b contains infs")
|
|
|
|
a = a.detach().cpu().numpy().flatten()
|
|
b = b.detach().cpu().numpy().flatten()
|
|
|
|
# compare elements of a and b relative to the max value in b
|
|
# large fp32 values may cause quantization errors that propagate to small values
|
|
rel_diff = np.abs(a-b)/np.linalg.norm(b)
|
|
abs_diff = np.abs(a-b)
|
|
cos_diff = distance.cosine(a, b)
|
|
try:
|
|
if rel_diff.max() > rtol:
|
|
raise AssertionError("Tensor relative error > %.2e (%.2e)" % (rtol, rel_diff.max()))
|
|
if abs_diff.max() > atol:
|
|
raise AssertionError("Tensor absolute error > %.2e (%.2e)" % (atol, abs_diff.max()))
|
|
if cos_diff > ctol:
|
|
raise AssertionError("Tensor cosine distance > %.2e (%.2e)" % (ctol, cos_diff))
|
|
# np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
|
|
# np.testing.assert_array_almost_equal_nulp(a, b)
|
|
except AssertionError as e:
|
|
print('norm(a) =', np.linalg.norm(a))
|
|
print('norm(b) =', np.linalg.norm(b))
|
|
print('Largest relative difference = %.2e' % rel_diff.max())
|
|
idx = np.argmax(rel_diff)
|
|
print('a[%d] = %.10f' % (idx, a[idx]))
|
|
print('b[%d] = %.10f' % (idx, b[idx]))
|
|
print('Largest absolute difference = %.2e' % abs_diff.max())
|
|
idx = np.argmax(abs_diff)
|
|
print('a[%d] = %.10f' % (idx, a[idx]))
|
|
print('b[%d] = %.10f' % (idx, b[idx]))
|
|
print('Cosine distance = %.2e' % cos_diff)
|
|
raise e
|
|
|
|
def assert_min_mse(a, b, tol=1e-20):
|
|
"""Assert that the mean squared error between a and b is at least tol."""
|
|
a = a.detach().cpu().numpy()
|
|
b = b.detach().cpu().numpy()
|
|
mse = ((a-b)**2).mean()
|
|
if mse < tol:
|
|
raise AssertionError("MSE = %.2e < %.2e" % (mse, tol))
|
|
|
|
def quant_np(x, amax, num_bits=8, fake=False, narrow_range=True):
|
|
"""Quantize x using numpy."""
|
|
intmax = 2.0**(num_bits - 1) - 1
|
|
intmin = -intmax if narrow_range else -intmax - 1
|
|
scale = intmax / amax
|
|
x_q = np.round(np.clip(x * scale, intmin, intmax))
|
|
|
|
if fake:
|
|
x_q /= scale
|
|
|
|
return x_q
|