chore: import upstream snapshot with attribution

tools/pytorch-quantization/tests/utils.py

@@ -0,0 +1,129 @@
+#
+# 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