chore: import upstream snapshot with attribution

tools/pytorch-quantization/tests/calibrator_test.py

@@ -0,0 +1,452 @@
+#
+# 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 calibrators"""
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization import calib
+from pytorch_quantization import nn as quant_nn
+import tests.utils as test_utils
+from tests.fixtures import verbose
+from tests.fixtures.models import QuantLeNet
+
+np.random.seed(12345)
+torch.manual_seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestMaxCalibrator():
+
+    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_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_reverse_axis(self):
+        axis = -4
+        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_raises(self):
+        axis = 0
+        max_calibrator = calib.MaxCalibrator(8, axis, False)
+
+        x_2 = torch.rand(32, 63, 7, 7).cuda()
+        x_3 = torch.rand(33, 63, 7, 7).cuda()
+        max_calibrator.collect(x_2)
+        with pytest.raises(RuntimeError, match="shape changed"):
+            max_calibrator.collect(x_3)
+
+    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_repr(self):
+        max_calibrator = calib.MaxCalibrator(8, None, False, track_amax=True)
+        repr(max_calibrator)
+
+class TestHistogramCalibrator():
+
+    def test_grow(self, verbose):
+        x_1 = torch.tensor([0, 255, 255, 255, 255, 255]).cuda()
+        x_2 = torch.tensor([0, 255, 255, 255, 255, 256]).cuda()
+
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='stretch')
+        hist_calibrator.collect(x_1)
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax(method='entropy')
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 256 because the last bin gets stretched to (~255, 257)
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='append')
+        hist_calibrator.collect(x_1)
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax(method='mse')
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < 0.5
+
+    def test_skip_zeros(self, verbose):
+        x_1 = torch.tensor([0, 0, 0, 0, 0, 1, 2, 3, 4, 5])
+        x_2 = torch.tensor([0, 0, 0, 0, 0, 6, 7, 8, 9, 10])
+
+        calibrator = calib.HistogramCalibrator(8, None, False, skip_zeros=True)
+        calibrator.collect(x_1)
+        calibrator.collect(x_2)
+
+        amax = calibrator.compute_amax("percentile", percentile=50)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be close to 5
+        assert (amax - 5.).abs() < 10/2048
+
+    def test_torch_hist(self):
+        x_1 = torch.rand(1023, device="cuda")
+        x_1[0] = 0
+        x_2 = torch.rand(1023, device="cuda") + 1  # Make sure histogram bins need to be grown
+        x_2[1] = 0
+
+        calibrator_np = calib.HistogramCalibrator(8, None, False, num_bins=19, torch_hist=False)
+        calibrator_torch = calib.HistogramCalibrator(8, None, False, num_bins=19, torch_hist=True)
+
+        calibrator_np.collect(x_1)
+        calibrator_torch.collect(x_1)
+        assert calibrator_torch._calib_hist.numel() == calibrator_torch._calib_bin_edges.numel() - 1
+        np.testing.assert_array_equal(calibrator_np._calib_hist, calibrator_torch._calib_hist.cpu().numpy())
+        np.testing.assert_array_almost_equal(
+            calibrator_np._calib_bin_edges, calibrator_torch._calib_bin_edges.cpu().numpy())
+
+        # Test multiple collections with some of them needs to expand range
+        for _ in range(3):
+            calibrator_np.collect(x_2)
+            calibrator_torch.collect(x_2)
+            calibrator_np.collect(x_1)
+            calibrator_torch.collect(x_1)
+
+            # Test compute_amax function doesn't convert _calib_hist and _calib_bin_edges unnecessarily
+            calibrator_np.compute_amax("percentile", percentile=99.99)
+            calibrator_torch.compute_amax("percentile", percentile=99.99)
+
+            np.testing.assert_array_equal(calibrator_np._calib_hist, calibrator_torch._calib_hist.cpu().numpy())
+            np.testing.assert_array_almost_equal(
+                calibrator_np._calib_bin_edges, calibrator_torch._calib_bin_edges.cpu().numpy())
+            assert calibrator_torch._calib_hist.numel() == calibrator_torch._calib_bin_edges.numel() - 1
+
+
+class TestEntropyCalibrator():
+
+    def test_one_tensor(self, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='stretch')
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+
+        # Don't have a better test metric. One outlier 10 should be discared by KL-divergence
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    def test_unsigned(self, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, True, grow_method='stretch')
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+        x_3 = torch.rand(11, 7, 3, 3).cuda()
+        hist_calibrator.collect(x_3)
+
+        # Don't have a better test metric. One outlier 10 should be discared by KL-divergence
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    def test_repr(self):
+        hist_calibrator = calib.HistogramCalibrator(8, None, True)
+        repr(hist_calibrator)
+
+class TestMSECalibrator():
+
+    def test_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        x_1[1, 1, 1, 1] = 256. # create an outlier
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+    def test_unsigned_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, True)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 512.
+        x_1[1, 1, 1, 1] = 513. # create an outlier
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 512
+        assert (amax - 512.).abs() < (amax - 513.).abs()
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        x_1[1, 1, 1, 1] = 256. # create an outlier
+        calibrator.collect(x_1)
+        x_2 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        calibrator.collect(x_2)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+    def test_repr(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        repr(calibrator)
+
+class TestPercentileCalibrator():
+
+    def test_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("percentile", percentile=90)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 89
+        assert (amax - 89.).abs() < 100/1024
+
+    def test_unsigned_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator( 8, None, True)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("percentile", percentile=80)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 79
+        assert (amax - 79.).abs() < 100/2048
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+        x_2 = torch.arange(0, 50, 0.5)
+        calibrator.collect(x_2)
+        amax = calibrator.compute_amax("percentile", percentile=99)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 97
+        assert (amax - 97.).abs() < 100/1024
+
+    def test_repr(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        repr(calibrator)
+
+    def test_range(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+        with pytest.raises(ValueError, match="range"):
+            calibrator.compute_amax("percentile", percentile=-10)
+        with pytest.raises(ValueError, match="range"):
+            calibrator.compute_amax("percentile", percentile=200)
+
+class TestCalibrateWeights():
+
+    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_shape_with_axis(self):
+        """Check calibrate_weight function returns same shape as TensorQuantizer"""
+        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="percentile")
+
+        for ref_module, test_module in zip(ref_lenet.modules(), test_lenet.modules()):
+            if isinstance(ref_module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
+                assert ref_module.weight_quantizer.amax.shape == test_module.weight_quantizer.amax.shape
+
+    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_percentile_with_axis(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=True, percentile=test_percentile)
+        ref_calibrator.collect(test_lenet.conv2.weight[1])
+        ref_amax = ref_calibrator.compute_amax("percentile", percentile=test_percentile)
+        test_utils.compare(ref_amax, test_lenet.conv2.weight_quantizer.amax[1], rtol=0, atol=0, ctol=0)
+
+    def test_mse(self):
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+
+        ref_calibrator = calib.HistogramCalibrator(8, None, False)
+
+        calib.calibrate_weights(test_lenet, method="mse", perchannel=False)
+        ref_calibrator.collect(test_lenet.conv1.weight)
+        ref_amax = ref_calibrator.compute_amax("mse")
+        test_utils.compare(ref_amax, test_lenet.conv1.weight_quantizer.amax, rtol=0, atol=0, ctol=0)
+
+    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)

tools/pytorch-quantization/tests/classification_flow_test.py

@@ -0,0 +1,72 @@
+#
+# 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 the classification flow"""
+import os
+import subprocess
+import sys
+from os import path
+import glob
+import pytest
+# pylint:disable=missing-docstring, no-self-use
+
+class TestClassificationFlow():
+
+    def test_resnet18(self, request, pytestconfig):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        dataset_dir = pytestconfig.getoption('--data-dir')
+
+        # skip if the data dir flag was not set
+        if not dataset_dir:
+            pytest.skip("Prepare required dataset and use --data-dir option to enable")
+
+        # Verify data dir exists
+        if not path.exists(dataset_dir):
+            print("Dataset path %s doesn't exist"%(dataset_dir), file=sys.stderr)
+            assert path.exists(dataset_dir)
+
+        # Append required paths to PYTHONPATH
+        test_env = os.environ.copy()
+        if 'PYTHONPATH' not in test_env:
+            test_env['PYTHONPATH'] = ""
+
+        # Add project root and torchvision to the path (assuming running in nvcr.io/nvidia/pytorch:20.08-py3)
+        test_env['PYTHONPATH'] += ":/opt/pytorch/vision/references/classification/:%s/../"%(dir_path)
+
+        # Add requirement egg files manually to path since we're spawning a new process (downloaded by setuptools)
+        for egg in glob.glob(dir_path + "/../.eggs/*.egg"):
+            test_env['PYTHONPATH'] += ":%s"%(egg)
+
+        # Run in a subprocess to avoid contaminating the module state for other test cases
+        ret = subprocess.run(
+            [
+                'python3', dir_path + '/../examples/torchvision/classification_flow.py',
+                '--data-dir', dataset_dir,
+                '--model', 'resnet18', '--pretrained',
+                '-t', '0.5',
+                '--num-finetune-epochs', '2',
+                '--evaluate-onnx',
+            ],
+            env=test_env,
+            check=False, stdout=subprocess.PIPE)
+
+        # If the test failed dump the output to stderr for better logging
+        if ret.returncode != 0:
+            print(ret.stdout, file=sys.stderr)
+
+        assert ret.returncode == 0

tools/pytorch-quantization/tests/clip_test.py

@@ -0,0 +1,69 @@
+#
+# 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 Clip module."""
+
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization.nn.modules import clip
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestClip():
+
+    def test_simple_run(self):
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np)
+        clip_op = clip.Clip(torch.tensor(0.3), torch.tensor(0.7))
+        clip_x_np = np.clip(x_np, 0.3, 0.7)
+        clip_x_torch = clip_op(x_torch)
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_raise(self):
+        with pytest.raises(ValueError, match="must be scalar"):
+            clip_op = clip.Clip(torch.tensor(0.3), torch.tensor(0.7), learn_min=True)
+
+    def test_backward(self):
+        x = torch.randn(3, 7, requires_grad=True)
+        x.retain_grad()
+
+        min_value = 0.3
+        max_value = 0.7
+        clip_op = clip.Clip(min_value, max_value, learn_min=True, learn_max=True)
+        clip_x = clip_op(x)
+        clip_x.retain_grad()
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+
+        loss.backward()
+
+        assert x.grad.cpu()[x.cpu() < min_value].sum() == 0
+        assert x.grad.cpu()[x.cpu() > max_value].sum() == 0
+        assert torch.equal(clip_x.grad[(x > min_value) & (x < max_value)], x.grad[(x > min_value) & (x < max_value)])

tools/pytorch-quantization/tests/conftest.py

@@ -0,0 +1,25 @@
+#
+# 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.
+#
+
+
+"""local configuration for pytests"""
+
+import pytest
+
+def pytest_addoption(parser):
+    parser.addoption('--data-dir', type=str, dest="data_dir",
+                     default='', help="set dataset dir for tests")

tools/pytorch-quantization/tests/fixtures/__init__.py

@@ -0,0 +1,22 @@
+#
+# 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.
+#
+
+
+import pytest
+@pytest.fixture
+def verbose(request):
+    return request.config.getoption("verbose")

tools/pytorch-quantization/tests/fixtures/models.py

@@ -0,0 +1,69 @@
+#
+# 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.
+#
+
+
+"""Model used for tests"""
+
+import pytest
+import torch.nn as nn
+import torch.nn.functional as F
+from pytorch_quantization.nn import QuantConv2d, QuantLinear
+from pytorch_quantization.tensor_quant import QuantDescriptor
+
+class LeNet(nn.Module):
+    def __init__(self, **kwargs):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5, **kwargs)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5, **kwargs)
+        self.fc1 = nn.Linear(320, 50, **kwargs)
+        self.fc2 = nn.Linear(50, 10, **kwargs)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2(x), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class QuantLeNet(nn.Module):
+    def __init__(self, **kwargs):
+        super(QuantLeNet, self).__init__()
+        self.conv1 = QuantConv2d(1, 10, kernel_size=5, **kwargs)
+        self.conv2 = QuantConv2d(10, 20, kernel_size=5, **kwargs)
+        self.fc1 = QuantLinear(320, 50, **kwargs)
+        self.fc2 = QuantLinear(50, 10, **kwargs)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2(x), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+@pytest.fixture
+def resnet18():
+    import torchvision
+    return torchvision.models.resnet18()
+
+@pytest.fixture
+def quant_lenet():
+    return QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor())

tools/pytorch-quantization/tests/functional_test.py

@@ -0,0 +1,94 @@
+#
+# 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 supportive functions"""
+
+import pytest
+import numpy as np
+
+import torch
+
+import pytorch_quantization.nn.functional as QF
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+
+class TestClip():
+
+    def test_simple_run(self):
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np)
+        clip_x_np = np.clip(x_np, 0.3, 0.7)
+        clip_x_torch = QF.clip(x_torch, torch.tensor(0.3), torch.tensor(0.7))
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_raise(self):
+        x = torch.randn(3, 7, requires_grad=True)
+
+        min_value = torch.Tensor(3, 7)
+        max_value = torch.Tensor(3, 7)
+        min_value.requires_grad = True
+        max_value.requires_grad = True
+        clip_x = QF.clip(x, min_value, max_value)
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+        with pytest.raises(ValueError, match="can only be scalar"):
+            loss.backward()
+
+    def test_broadcast(self):
+        """Test broadcast behavior by randomly picked shuffling of np.random.rand"""
+        x_np = np.random.rand(1023, 4, 5, 6).astype(np.float32) - 0.5
+        x_torch = torch.Tensor(x_np)
+        min_value = np.random.rand(1, 4, 1, 1).astype(np.float32) * 0.1 - 0.2
+        max_value = np.random.rand(1, 4, 1, 1).astype(np.float32) * 10 + 0.5
+        clip_x_np = np.clip(x_np, min_value, max_value)
+        clip_x_torch = QF.clip(x_torch, torch.tensor(min_value), torch.tensor(max_value))
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_backward(self):
+        x = torch.randn(3, 1025, requires_grad=True)
+        x.retain_grad()
+
+        min_value = torch.tensor(0.3)
+        max_value = torch.tensor(0.7)
+        min_value.requires_grad = True
+        max_value.requires_grad = True
+        min_value.retain_grad()
+        max_value.retain_grad()
+        clip_x = QF.clip(x, min_value, max_value)
+        clip_x.retain_grad()
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+        loss.backward()
+
+        np.testing.assert_array_almost_equal(
+            clip_x.grad[x < min_value].sum().cpu().numpy(), min_value.grad.cpu().numpy(), decimal=6)
+        np.testing.assert_array_almost_equal(
+            clip_x.grad[x > max_value].sum().cpu().numpy(), max_value.grad.cpu().numpy(), decimal=6)
+        assert x.grad.cpu()[x.cpu() < min_value.cpu()].sum() == 0
+        assert x.grad.cpu()[x.cpu() > max_value.cpu()].sum() == 0
+        assert torch.equal(clip_x.grad[(x > min_value) & (x < max_value)], x.grad[(x > min_value) & (x < max_value)])

tools/pytorch-quantization/tests/integration_test.py

@@ -0,0 +1,216 @@
+#
+# 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 integrating Quant layers into a network"""
+
+import pytest
+import io
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization import quant_modules
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules import tensor_quantizer
+from tests.fixtures.models import LeNet, QuantLeNet
+from tests.fixtures import verbose
+
+np.random.seed(12345)  # seed 1234 causes 1 number mismatch at 6th decimal in one of the tests
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestNetwork():
+    """test basic operations of quantized network"""
+
+    def test_simple_build(self):
+        """test instantiation"""
+        quant_model = QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor())
+        for name, module in quant_model.named_modules():
+            if "quantizer" in name:
+                module.disable()
+
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+
+        input_desc = QuantDescriptor(amax=6.)
+        weight_desc = QuantDescriptor(amax=1.)
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+
+
+    def test_forward(self):
+        """test forward pass with random data"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        output = quant_model(torch.empty(16, 1, 28, 28))
+
+    def test_backward(self):
+        """test one iteration with random data and labels"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        optimizer = optim.SGD(quant_model.parameters(), lr=0.01)
+        optimizer.zero_grad()
+        output = quant_model(torch.empty(16, 1, 28, 28))
+        loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+
+    def test_native_amp_fp16(self):
+        """test one iteration with random data and labels"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        optimizer = optim.SGD(model.parameters(), lr=0.01)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast():
+            output = model(torch.empty(16, 1, 28, 28))
+            loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+        assert loss.dtype == torch.float32
+
+    def test_asp(self):
+        """test Sparsity (ASP) and QAT toolkits together"""
+        try:
+            from apex.contrib.sparsity import ASP
+        except ImportError:
+            pytest.skip("ASP is not available.")
+
+        quant_modules.initialize()
+        model = LeNet()
+        quant_modules.deactivate()
+
+        optimizer = optim.SGD(model.parameters(), lr=0.01)
+
+        ASP.init_model_for_pruning(
+            model,
+            mask_calculator="m4n2_1d",
+            verbosity=2,
+            whitelist=[torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv3d, quant_nn.modules.quant_linear.QuantLinear],
+            allow_recompute_mask=False,
+            custom_layer_dict={
+                quant_nn.QuantConv1d: ['weight'],
+                quant_nn.QuantConv2d: ['weight'],
+                quant_nn.QuantConv3d: ['weight'],
+                quant_nn.QuantConvTranspose1d: ['weight'],
+                quant_nn.QuantConvTranspose2d: ['weight'],
+                quant_nn.QuantConvTranspose3d: ['weight'],
+                quant_nn.QuantLinear: ['weight']
+            },
+            allow_permutation=False)
+        ASP.init_optimizer_for_pruning(optimizer)
+        ASP.compute_sparse_masks()
+
+        model = model.to('cuda')
+        output = model(torch.empty(16, 1, 28, 28).to('cuda'))
+        optimizer.zero_grad()
+        loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+
+    def test_quant_module_replacement(self):
+        """test monkey patching of modules with their quantized versions"""
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]  
+
+        # Before any monkey patching, the networks should be different
+        assert(mod_list != qmod_list)
+
+        # Monkey patch the modules
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+
+        quant_modules.initialize(no_replace_list, custom_quant_modules)
+
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+    
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]
+
+        # After monkey patching, the networks should be same
+        assert(mod_list == qmod_list)
+
+        # Reverse monkey patching
+        quant_modules.deactivate()
+
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+    
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]
+
+        # After reversing monkey patching, the networks should again be different
+        assert(mod_list != qmod_list)
+
+    def test_calibration(self):
+        quant_model = QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor()).cuda()
+
+        for name, module in quant_model.named_modules():
+            if name.endswith("_quantizer"):
+                if module._calibrator is not None:
+                    module.disable_quant()
+                    module.enable_calib()
+                else:
+                    module.disable()
+                print(F"{name:40}: {module}")
+
+        quant_model(torch.rand(16, 1, 224, 224, device="cuda"))
+
+        # Load calib result and disable calibration
+        for name, module in quant_model.named_modules():
+            if name.endswith("_quantizer"):
+                if module._calibrator is not None:
+                    module.load_calib_amax()
+                    module.enable_quant()
+                    module.disable_calib()
+                else:
+                    module.enable()
+        quant_model.cuda()
+
+    def test_state_load(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()
+        quantizer2.pre_quant_scale = torch.Tensor([[1.0, 2.0, 3.0, 4.0]]).cuda()
+        buffer = io.BytesIO()
+        torch.save(quantizer2.state_dict(), buffer)
+
+        buffer.seek(0)
+        quantizer.load_state_dict(torch.load(buffer))
+
+        assert torch.allclose(quantizer.pre_quant_scale, quantizer2.pre_quant_scale)

tools/pytorch-quantization/tests/license_test.py

@@ -0,0 +1,86 @@
+#
+# 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.
+#
+
+
+"""test the license of source files."""
+
+import pytest
+from pathlib import Path
+from filecmp import cmp
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestLicense():
+
+    def test_license(self):
+        root = Path(__file__).parent.parent.absolute()
+        root_len = len(str(root))
+
+        # Collect files ending with relevant extensions
+        file_list = []
+        file_types = ['*.py', '*.cpp', '*.cu', '*.h', '*.hpp', '*.c', '*.sh']
+        for ft in file_types:
+            file_list += list(root.rglob(ft))
+
+        # Trim files from build folders
+        build_folders = ['build', 'dist', '.eggs', '.vscode']
+        build_files = []
+        for src_file in file_list:
+            local_path = str(src_file.parents[0])[root_len : ]
+            for folder in build_folders:
+                if folder in local_path:
+                    build_files.append(src_file)
+
+        for bf in build_files:
+            file_list.remove(bf)
+
+        print (f"Found {len(file_list)} source files")
+
+        cpp_header = (root / 'tests' / 'license_test_header_cpp.txt').open().readlines()
+        py_header = (root / 'tests' / 'license_test_header_py.txt').open().readlines()
+        sh_header = (root / 'tests' / 'license_test_header_sh.txt').open().readlines()
+
+        invalid_files = []
+        for f in file_list:
+            with open(f) as src_file:
+                src_lines = src_file.readlines()
+
+            # Skip empty files
+            if len(src_lines) == 0:
+                continue
+
+            if f.suffix == '.py':
+                header = py_header
+            elif f.suffix == '.sh':
+                header = sh_header
+            else:
+                header = cpp_header
+
+            num_lines = len(header)
+            if len(src_lines) < num_lines:
+                invalid_files.append(f)
+                continue
+
+            for i in range(num_lines):
+                if src_lines[i] != header[i]:
+                    invalid_files.append(f)
+                    break
+
+        if len(invalid_files) > 0:
+            for f in invalid_files:
+                print(f"The file {f} has an invalid header!")
+            raise AssertionError("%d files have invalid headers!" % (len(invalid_files)))

tools/pytorch-quantization/tests/license_test_header_cpp.txt

@@ -0,0 +1,16 @@
+/*
+ * 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.
+ */

tools/pytorch-quantization/tests/license_test_header_py.txt

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

tools/pytorch-quantization/tests/license_test_header_sh.txt

@@ -0,0 +1,17 @@
+#!/bin/bash
+#
+# 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.
+#

tools/pytorch-quantization/tests/model_test.py

@@ -0,0 +1,92 @@
+#
+# 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 calibrators"""
+import inspect
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import enable_onnx_export
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization import calib
+from pytorch_quantization import nn as quant_nn
+import tests.utils as test_utils
+from examples.torchvision.models.classification import *
+from tests.fixtures import verbose
+from tests.fixtures.models import QuantLeNet
+
+np.random.seed(12345)
+torch.manual_seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestExampleModels():
+
+    def test_resnet50(self):
+        model = resnet50(pretrained=True, quantize=True)
+        model.eval()
+
+        for name, module in model.named_modules():
+            if name.endswith('_quantizer'):
+                module.amax = 2.50
+
+        model.cuda()
+        dummy_input = torch.randn(1, 3, 224, 224, device='cuda')
+        with enable_onnx_export():
+            if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  enable_onnx_checker=False,
+                                  do_constant_folding=True)
+            else:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  do_constant_folding=True)
+
+    def test_resnet50_cpu(self):
+        model = resnet50(pretrained=True, quantize=True)
+        model.eval()
+
+        for name, module in model.named_modules():
+            if name.endswith('_quantizer'):
+                module.amax = 2.50
+
+        dummy_input = torch.randn(1, 3, 224, 224)
+        with enable_onnx_export():
+            if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50_cpu.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  enable_onnx_checker=False,
+                                  do_constant_folding=True)
+            else:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  do_constant_folding=True)

tools/pytorch-quantization/tests/optim_helper_test.py

@@ -0,0 +1,111 @@
+#
+# 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 helper functions for quant optimizer"""
+
+import numpy as np
+import pytest
+
+
+import torch.optim as optim
+
+from pytorch_quantization.optim import helper
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from .fixtures.models import QuantLeNet
+from .fixtures.models import resnet18
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestMatchParameters():
+
+    def test_single_key(self, resnet18):
+        param = helper.match_parameters(resnet18, ['downsample.0.weight'])
+        assert len(list(param)) == 3
+
+    def test_multi_keys(self, resnet18):
+        param = list(helper.match_parameters(resnet18, ['conv1', 'downsample']))
+        assert len(param) == 18
+
+    def test_regex(self, resnet18):
+        param = helper.match_parameters(resnet18, ['downsample.*.weight$'])
+        assert len(list(param)) == 6
+
+        param = helper.match_parameters(resnet18, ['downsample.*.wei$'])
+        assert not list(param)
+
+
+class TestGroupParameters():
+
+    def test_single_key(self, resnet18):
+        param_groups = helper.group_parameters(resnet18, [['downsample.1.weight']])
+        assert len(list(param_groups[0]['params'])) == 3
+
+    def test_lr_momentum_decay(self, resnet18):
+        lrs = [0.01, 0.001]
+        momentums = [0.02, 0.002]
+        weight_decays = [0.03, 0.003]
+        param_groups = helper.group_parameters(
+            resnet18, [['conv1.*weight'], ['downsample.*.weight']], lrs, momentums, weight_decays)
+
+        assert param_groups[0]['lr'] == lrs[0]
+        assert param_groups[1]['lr'] == lrs[1]
+        assert param_groups[0]['momentum'] == momentums[0]
+        assert param_groups[1]['momentum'] == momentums[1]
+        assert param_groups[0]['weight_decay'] == weight_decays[0]
+        assert param_groups[1]['weight_decay'] == weight_decays[1]
+
+    def test_optimizer_feed(self, resnet18):
+        """Feed grouped parameters to optimizer, see what happens"""
+        lrs = [0.01, 0.001]
+        momentums = [0.02, 0.002]
+        weight_decays = [0.03, 0.003]
+        param_groups = helper.group_parameters(
+            resnet18, [['conv1.*weight'], ['downsample.*.weight']], lrs, momentums, weight_decays)
+        optimizer = optim.SGD(param_groups)
+        optimizer.step()
+
+    def test_raises(self):
+        with pytest.raises(TypeError, match="must be list of list of patterns"):
+            helper.group_parameters(None, [['downsample.1.weight'], 'conv1'])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], lrs=[0.1])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], momentums=[0.1])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], weight_decays=[0.1])
+
+
+class TestFreezeParameters():
+
+    def test_simple(self, resnet18):
+        helper.freeze_parameters(resnet18, ['downsample.0.weight'])
+        for name, param in resnet18.named_parameters():
+            if 'downsample.0.weight' in name:
+                assert not param.requires_grad
+
+class TestQuantWeightInPlace():
+
+    def test_simple(self):
+        quant_lenet = QuantLeNet(
+            quant_desc_input=QuantDescriptor(),
+            quant_desc_weight=QuantDescriptor())
+        quant_lenet.eval()
+        helper.quant_weight_inplace(quant_lenet)

tools/pytorch-quantization/tests/print_test.py

@@ -0,0 +1,60 @@
+#
+# 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.
+#
+
+
+"""test for str and repr
+Make sure things can print and in a nice form. Put all the print tests together so that running this test file alone
+can inspect all the print messages in the project
+"""
+
+import torch
+from torch import nn
+
+from pytorch_quantization import calib
+from pytorch_quantization import tensor_quant
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestPrint():
+
+    def test_print_descriptor(self):
+        test_desc = tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL
+        print(test_desc)
+
+    def test_print_tensor_quantizer(self):
+        test_quantizer = TensorQuantizer()
+        print(test_quantizer)
+
+    def test_print_module(self):
+        class _TestModule(nn.Module):
+            def __init__(self):
+                super(_TestModule, self).__init__()
+                self.conv = nn.Conv2d(33, 65, 3)
+                self.quant_conv = quant_nn.Conv2d(33, 65, 3)
+                self.linear = nn.Linear(33, 65)
+                self.quant_linear = quant_nn.Linear(33, 65)
+
+        test_module = _TestModule()
+        print(test_module)
+
+    def test_print_calibrator(self):
+        print(calib.MaxCalibrator(7, 1, False))
+        hist_calibrator = calib.HistogramCalibrator(8, None, True)
+        hist_calibrator.collect(torch.rand(10))
+        print(hist_calibrator)

tools/pytorch-quantization/tests/quant_conv_test.py

@@ -0,0 +1,557 @@
+#
+# 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 QuantConv module.
+Mose tests check the functionality of all the combinations in Quant conv against the corresponding functionalities in
+tensor_quant. There are tests for all the three QuantConv1D, QuantConv2D, and QuantConv3D
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_conv
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+_NUM_IN_CHANNELS = 13
+_NUM_OUT_CHANNELS = 17
+
+class TestQuantConv2D():
+    #Quantizing weight
+
+    def test_no_quant(self):
+
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_in_feature_fake_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv2d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False,
+                                                   quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3)
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True,
+                                                   quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24, 24).cuda()
+
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        fake_quant_conv2d = quant_conv.QuantConv2d(
+            _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(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        conv2d = nn.Conv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv2d(test_input)
+        output = conv2d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-6, atol=1.5e-4)
+
+
+    def test_set_default_quant_desc(self):
+        quant_conv_layer = quant_conv.Conv2d(32, 257, 3)
+        assert quant_conv_layer.input_quantizer._axis == None
+        assert quant_conv_layer.weight_quantizer._axis == (0)
+
+        # set default to a different one
+        quant_desc_input = QuantDescriptor(num_bits=11)
+        quant_desc_weight = QuantDescriptor(num_bits=13, axis=(1))
+        quant_conv.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+        quant_conv.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
+
+        # Create one with default descriptor
+        quant_conv_layer = quant_conv.Conv2d(32, 257, 3)
+        # Check quant_desc in quantizer created with default descriptor
+        assert quant_conv_layer.input_quantizer._num_bits == quant_desc_input.num_bits
+        assert quant_conv_layer.weight_quantizer._axis == quant_desc_weight.axis
+
+        # Test default is per class
+        quant_conv_layer = quant_conv.Conv3d(31, 255, 5)
+        assert quant_conv_layer.input_quantizer._num_bits != quant_desc_input.num_bits
+        assert quant_conv_layer.weight_quantizer._axis != quant_desc_weight.axis
+
+        # Reset default
+        quant_conv.QuantConv2d.set_default_quant_desc_input(QuantDescriptor())
+        quant_conv.QuantConv2d.set_default_quant_desc_weight(QuantDescriptor(axis=(0)))
+
+    def test_unused_kwargs(self):
+        with pytest.raises(TypeError, match="Unused keys"):
+            quant_conv.Conv2d(32, 257, 3, descriptor='oops')
+
+class TestQuantConv1D():
+
+    def test_no_quant(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(1, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv1d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1))
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(0))
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1))
+
+        out1 = F.conv1d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24).cuda()
+
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        fake_quant_conv1d = quant_conv.QuantConv1d(
+            _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(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        conv1d = nn.Conv1d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv1d(test_input)
+        output = conv1d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=1e-4)
+
+
+class TestQuantConv3D():
+    #Quantizing weight
+
+    def test_no_quant(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv3d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(0))
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv3d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(8, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1, 1))
+
+        out1 = F.conv3d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    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)

tools/pytorch-quantization/tests/quant_conv_transposed_test.py

@@ -0,0 +1,522 @@
+#
+# 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 QuantConv module.
+Test for QuantConvTransposed
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_conv
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+_NUM_IN_CHANNELS = 13
+_NUM_OUT_CHANNELS = 17
+
+
+class TestQuantConvTranspose2D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(256, _NUM_IN_CHANNELS, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv_transpose2d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv2d = quant_conv.QuantConvTranspose2d(
+            _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=(1)))
+
+        # 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)
+        conv2d = nn.ConvTranspose2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv2d(test_input)
+        output = conv2d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=2e-4)
+
+
+class TestQuantConvTranspose3D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose3d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose3d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2, 2, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3, 4))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose3d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32, 32).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv3d = quant_conv.QuantConvTranspose3d(
+            _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=(1)))
+
+        # 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.ConvTranspose3d(_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-5, atol=2e-4)
+
+
+class TestQuantConvTranspose1D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(256, _NUM_IN_CHANNELS, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv_transpose1d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv1d = quant_conv.QuantConvTranspose1d(
+            _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=(1)))
+
+        # 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)
+        conv1d = nn.ConvTranspose1d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv1d(test_input)
+        output = conv1d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=1e-4)

tools/pytorch-quantization/tests/quant_instancenorm_test.py

@@ -0,0 +1,199 @@
+#
+# 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 QuantInstanceNorm module.
+Mose tests check the functionality of all the combinations in Quant instancenorm against the corresponding functionalities in
+tensor_quant. There are tests for all the three QuantInstaceNorm1D, QuantInstanceNorm2D, and QuantInstanceNorm3D
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_instancenorm
+#import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+NUM_CHANNELS = 15
+
+class TestQuantInstanceNorm1D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+
+class TestQuantInstanceNorm2D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0].max(3, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+
+
+class TestQuantInstanceNorm3D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0]
+                                                          .max(3, keepdim=True)[0].max(4, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())

tools/pytorch-quantization/tests/quant_linear_test.py

@@ -0,0 +1,231 @@
+#
+# 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 QuantLinear module.
+Most tests check the functionality of all the combinations in Quant Linear against the corresponding functionalities
+in tensor_quant.
+"""
+import pytest
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_linear
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestQuantLinear():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_linear_object = quant_linear.QuantLinear(
+                7, 9, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    #Quantizing weight
+    def test_weight_fake_per_tensor(self):
+        with torch.cuda.device(0):
+            size = 256
+            quant_linear_object = quant_linear.QuantLinear(
+                size,
+                size,
+                bias=False,
+                quant_desc_weight=tensor_quant.QuantDescriptor(axis=None))
+            quant_linear_object.input_quantizer.disable()
+            test_input = torch.randn(size, size)
+
+            weight_copy = quant_linear_object.weight.clone()
+            quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+            out1 = F.linear(test_input, quant_weight)
+            out2 = quant_linear_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_per_channel(self):
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
+        quant_linear_object.input_quantizer.disable()
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=1, keepdims=True)
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.linear(test_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    # Quantizing activations
+    def test_test_input_fake_per_tensor(self):
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False)
+        quant_linear_object.weight_quantizer.disable()
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.linear(quant_input, weight_copy)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor())
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor_with_bias(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor())
+        test_input = torch.randn(32, 17, 93, size_in)  # Test input other than 2 dimensional
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.linear(quant_input, quant_weight, bias=quant_linear_object.bias)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(size_in, size_out, bias=False,
+                                                       quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy,
+                                                      torch.max(torch.abs(weight_copy), dim=1, keepdim=True)[0])
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_precs(self):
+        """Test some precisions other than 8bit."""
+        size_in = 255
+        size_out = 257
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=4)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=3)
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in,
+            size_out,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = test_input_quantizer(test_input)
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_against_unquantized(self):
+        """
+        Quantized Linear should introduce bounded error compare to Linear
+        """
+        size_in = 255
+        size_out = 257
+        test_input = torch.randn(32, size_in).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        quant_linear_layer = quant_linear.QuantLinear(
+            size_in,
+            size_out,
+            bias=True,
+            quant_desc_input=tensor_quant.QuantDescriptor(num_bits=16),
+            quant_desc_weight=tensor_quant.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)
+        linear_layer = nn.Linear(size_in, size_out, bias=True)
+
+        quant_out_features = quant_linear_layer(test_input)
+        out_features = linear_layer(test_input)
+
+        # The difference between Linear and QuantLinear should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        np.testing.assert_allclose(
+            quant_out_features.detach().cpu().numpy(), out_features.detach().cpu().numpy(), rtol=0.01, atol=1e-4)
+
+    def test_set_default_quant_desc(self):
+        quant_linear_layer = quant_linear.QuantLinear(32, 257)
+        assert quant_linear_layer.input_quantizer.axis == None
+        assert quant_linear_layer.weight_quantizer.axis == (0)
+
+        # set default to a different one
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=11)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=13, axis=1)
+        quant_linear.Linear.set_default_quant_desc_input(quant_desc_input)
+        quant_linear.Linear.set_default_quant_desc_weight(quant_desc_weight)
+
+        # Create one with default descriptor
+        quant_linear_layer = quant_linear.QuantLinear(32, 257)
+        # Check quant_desc in quantizer created with default descriptor
+        assert quant_linear_layer.input_quantizer.num_bits == quant_desc_input.num_bits
+        assert quant_linear_layer.weight_quantizer.axis == quant_desc_weight.axis
+
+    def test_unused_kwargs(self):
+        with pytest.raises(TypeError, match="Unused keys"):
+            quant_linear_layer = quant_linear.QuantLinear(32, 257, descriptor='oops')

tools/pytorch-quantization/tests/quant_modules_test.py

@@ -0,0 +1,83 @@
+#
+# 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 Quant Module Replacement"""
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization import quant_modules
+from pytorch_quantization.quant_modules import QuantModuleReplacementHelper
+import tests.utils as test_utils
+from tests.fixtures import verbose
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestQuantModuleReplace():
+
+    def test_simple_default_args(self):
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state()
+        replacement_helper.apply_quant_modules()
+
+        # Linear module should not be replaced with its quantized version
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_with_no_replace_list(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = None
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state(no_replace_list, custom_quant_modules)
+        replacement_helper.apply_quant_modules()
+
+        # Linear module should not be replaced with its quantized version
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) != type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_with_custom_quant_modules(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state(no_replace_list, custom_quant_modules)
+        replacement_helper.apply_quant_modules()
+
+        # Although no replace list indicates Linear module should not be replaced with its
+        # quantized version, since the custom_quant_modules still contains the Linear module's
+        # mapping, it will replaced.
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_initialize_deactivate(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+
+        quant_modules.initialize(no_replace_list, custom_quant_modules)
+
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        quant_modules.deactivate()

tools/pytorch-quantization/tests/quant_pooling_test.py

@@ -0,0 +1,321 @@
+#
+# 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 QuantPooling module.
+Most tests check the functionality of all the combinations in Quant Pooling against the corresponding functionalities
+in tensor_quant.
+"""
+import pytest
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn.modules import quant_pooling
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+class TestQuantMaxPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool1d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantMaxPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.max_pool2d(test_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_multi_axis(self):
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8, axis=(0, 1))
+
+        quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(16, 7, 5, 5, dtype=torch.double)
+        input_amax = torch.amax(torch.abs(test_input), dim=(2, 3), keepdim=True)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, input_amax)
+
+        out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantMaxPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool3d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool1d(quant_input, 3, 1, 0, False, True)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantAvgPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.avg_pool2d(test_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool3d(quant_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool1d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool2d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantAdaptiveAvgPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.adaptive_avg_pool2d(quant_input, 3)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.adaptive_avg_pool2d(test_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool3d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())

tools/pytorch-quantization/tests/quant_rnn_test.py

@@ -0,0 +1,520 @@
+#
+# 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 QuantRNN module.
+"""
+import pytest
+
+import torch
+from torch import nn
+
+import numpy as np
+
+from pytorch_quantization.nn.modules import quant_rnn
+from pytorch_quantization import tensor_quant
+
+from tests.fixtures import verbose
+
+from . import utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+# change default type to double if utils.compare flags a small error, may just be floating point rounding error
+# torch.set_default_tensor_type('torch.cuda.DoubleTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(1234)
+
+# pylint: disable=no-self-use, missing-docstring, redefined-builtin, bad-continuation
+
+# global state for saving/loading test vectors
+SAVE_VECTORS = 0
+VECTOR_FILE = 'tests/quant_rnn_test_vectors.pt'
+if SAVE_VECTORS:
+    TEST_VECTORS = dict()
+else:
+    TEST_VECTORS = torch.load(VECTOR_FILE)
+
+
+class TestQuantLSTMCell():
+    """
+    tests for quant_rnn.QuantLSTMCell
+    default parameters in QuantLSTMCell:
+        bias=True,
+        num_bits_weight=8, quant_mode_weight='per_channel',
+        num_bits_input=8, quant_mode_input='per_tensor'
+
+    Tests of real quantization mode (nonfake) are disabled as it is not fully supported yet.
+    """
+
+    def test_basic_forward(self, verbose):
+        """Do a forward pass on the cell module and see if anything catches fire."""
+        batch = 7
+        input_size = 11
+        hidden_size = 9
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=8, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_rnn_object(input, hx=(hidden, cell))
+
+    def test_no_quant_input_hidden(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for input and hidden inputs."""
+        batch = 17
+        input_size = 13
+        hidden_size = 7
+
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_input_hidden_bias(self, verbose):
+        """QuantLSTMCell with quantization disabled vs. pytorch LSTMCell for input, hidden inputs and bias."""
+        batch = 19
+        input_size = 11
+        hidden_size = 3
+
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=True)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=True)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_against_unquantized(self, verbose):
+        """Quantization should introduce bounded error utils.compare to pytorch implementation."""
+        batch = 9
+        input_size = 13
+        hidden_size = 7
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=16)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=16, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        # The difference between reference and quantized should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        utils.compare(quant_hout, ref_hout, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_cout, ref_cout, rtol=1e-4, atol=1e-4)
+
+        # check that quantization introduces some error
+        utils.assert_min_mse(quant_hout, ref_hout, tol=1e-20)
+        utils.assert_min_mse(quant_cout, ref_cout, tol=1e-20)
+
+    def test_quant_input_hidden(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorchLSTMCell."""
+        batch = 15
+        input_size = 121
+        hidden_size = 51
+        num_bits = 4
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_input_hidden_bias(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorchLSTMCell
+            bias should not be quantized
+        """
+        batch = 9
+        input_size = 23
+        hidden_size = 31
+        num_bits = 7
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=True,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=True)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_different_prec(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorch LSTMCell
+            different input and weight precisions
+        """
+        batch = 27
+        input_size = 11
+        hidden_size = 10
+        num_bits_weight = 4
+        num_bits_input = 8
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits_input)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits_weight)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits_input)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits_weight)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+
+class TestQuantLSTM():
+    """
+    tests for quant_rnn.QuantLSTM
+    default parameters in QuantLSTM:
+        bias=True,
+        quant_weight=True, bits_weight=8, fake_quantTrue, quant_mode_weight='channel',
+        quant_input=True, bits_acts=8, quant_mode_input='tensor'
+
+    Tests of real quantization mode (nonfake) are disabled as it is not fully supported yet.
+    """
+
+    def test_basic_forward(self, verbose):
+        """Do a forward pass on the layer module and see if anything catches fire."""
+        batch = 5
+        input_size = 13
+        hidden_size = 31
+        seq_len = 1
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=8, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+        quant_rnn_object(input, hx=(hidden, cell))
+
+    def test_no_quant(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM."""
+        batch = 11
+        input_size = 14
+        hidden_size = 22
+        seq_len = 1
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+        quant_rnn_object._input_quantizers[0].disable()
+        quant_rnn_object._weight_quantizers[0].disable()
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input)
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input)
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_input_hidden(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for input and hidden inputs."""
+        batch = 13
+        input_size = 19
+        hidden_size = 20
+        seq_len = 1
+
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+        quant_rnn_object._input_quantizers[0].disable()
+        quant_rnn_object._weight_quantizers[0].disable()
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_all_modes(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for all modes."""
+
+        def testcase(input_size, hidden_size, seq_len, batch, num_layers, bias, batch_first, dropout, bidirectional):
+
+            quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                    num_layers=num_layers, bias=bias, batch_first=batch_first, dropout=dropout,
+                    bidirectional=bidirectional)
+
+            num_quantizers = num_layers * 2 if bidirectional else num_layers
+            for i in range(num_quantizers):
+                quant_rnn_object._input_quantizers[i].disable()
+                quant_rnn_object._weight_quantizers[i].disable()
+
+            ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                    num_layers=num_layers, bias=bias, batch_first=batch_first, dropout=dropout,
+                    bidirectional=bidirectional)
+
+            # copy state from one rnn to the other
+            ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+            input = torch.randn(seq_len, batch, input_size)
+            num_directions = 2 if bidirectional else 1
+            hidden = torch.randn(num_layers*num_directions, batch, hidden_size)
+            cell = torch.randn(num_layers*num_directions, batch, hidden_size)
+
+            quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+            ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+            utils.compare(quant_out, ref_out)
+            utils.compare(quant_hout, ref_hout)
+            utils.compare(quant_cout, ref_cout)
+
+        # test various permuatations of the following parameters:
+        #   size, num_layers, bias, batch_first, dropout, bidirectional
+        testcase(32, 27, 1, 1, 1, False, False, 0, False)
+        testcase(19, 63, 1, 1, 2, False, False, 0, False)
+        testcase(11, 41, 1, 1, 1, True, False, 0, False)
+        testcase(33, 31, 1, 1, 1, False, True, 0, False)
+        # testcase(32, 32, 1, 1, 2, False, False, 0.5, False) #TODO(pjudd) this fails look into dropout seeding
+        testcase(73, 13, 1, 1, 1, False, False, 0, True)
+
+    def test_against_unquantized(self, verbose):
+        """Quantization should introduce bounded error utils.compare to pytorch implementation."""
+        batch = 21
+        input_size = 33
+        hidden_size = 25
+        seq_len = 1
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=16)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=16, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+        # The difference between reference and quantized should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        utils.compare(quant_out, ref_out, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_hout, ref_hout, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_cout, ref_cout, rtol=1e-4, atol=1e-4)
+
+        # check that quantization introduces some error
+        utils.assert_min_mse(quant_out, ref_out, tol=1e-20)
+        utils.assert_min_mse(quant_hout, ref_hout, tol=1e-20)
+        utils.assert_min_mse(quant_cout, ref_cout, tol=1e-20)
+
+    def test_quant_input_hidden(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 13
+        input_size = 17
+        hidden_size = 7
+        seq_len = 1
+        num_bits = 6
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_input_hidden_bias(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 17
+        input_size = 13
+        hidden_size = 7
+        seq_len = 1
+        num_bits = 5
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=True,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=True,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_different_prec(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 22
+        input_size = 23
+        hidden_size = 24
+        seq_len = 1
+        num_bits_weight = 4
+        num_bits_input = 8
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits_input)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits_weight)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits_input)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits_weight)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+
+class TestEpilogue():
+    """Run after all tests to save globals."""
+
+    def test_save_vectors(self, verbose):
+        """Save test vectors to file."""
+        if SAVE_VECTORS:
+            torch.save(TEST_VECTORS, VECTOR_FILE)
+            raise Exception('Saved test vectors to {}, for testing set SAVE_VECTORS = 0'.format(VECTOR_FILE))

tools/pytorch-quantization/tests/quant_utils_test.py

@@ -0,0 +1,52 @@
+#
+# 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.
+#
+
+
+"""Test pytorch_quantization.utils"""
+import pytest
+import numpy as np
+
+import torch
+from pytorch_quantization import utils as quant_utils
+from tests.fixtures import verbose
+
+np.random.seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestQuantUtils():
+
+    def test_reduce_amax(self):
+        x_np = (np.random.rand(3, 7, 11, 13, 17) - 0.1).astype(np.float32)
+        x_torch = torch.tensor(x_np)
+
+        # Test reduce to one value
+        amax_np = np.max(np.abs(x_np))
+        amax_torch = quant_utils.reduce_amax(x_torch)
+        np.testing.assert_array_equal(amax_np, amax_torch.cpu().numpy())
+
+        # Test different axis
+        axes = [(1, 2, 3), (0, 2, 3), (0, 3), (0, 1, 3, 4)]
+        for axis in axes:
+            keepdims = np.random.rand() > 0.5
+            amax_np = np.max(np.abs(x_np), axis=axis, keepdims=keepdims)
+            amax_torch = quant_utils.reduce_amax(x_torch, axis=axis, keepdims=keepdims)
+            np.testing.assert_array_almost_equal(amax_np, amax_torch.cpu().numpy())
+
+        with pytest.raises(ValueError) as excinfo:
+            quant_utils.reduce_amax(x_torch, axis=(0, 1, 2, 3, 4, 5))
+            assert "Cannot reduce more axes" in str(excinfo.value)

tools/pytorch-quantization/tests/tensor_quant_test.py

@@ -0,0 +1,482 @@
+#
+# 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 quantization function and module"""
+import contextlib
+
+import pytest
+import numpy as np
+
+import torch
+from torch.nn.parameter import Parameter
+
+from pytorch_quantization import cuda_ext
+from pytorch_quantization import tensor_quant
+import pytorch_quantization.utils as quant_utils
+
+import tests.utils as test_utils
+from tests.fixtures import verbose
+
+np.random.seed(123456)  # seed 1234 causes 1 number mismatch at 6th decimal in one of the tests
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestTensorQuant():
+
+    def test_simple_run(self):
+        """ quantizer passes gradcheck
+        """
+        x = Parameter(torch.randn(2, 3, dtype=torch.float64).cuda()) * 100
+        tensor_quant.tensor_quant(x, torch.max(torch.abs(x)), 7)
+
+    def test_per_tensor_scale(self):
+        """ tensor_quant matches numpy quantization
+        """
+        torch.set_default_tensor_type('torch.cuda.FloatTensor')  # Test on GPU
+        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)))
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        torch.set_default_tensor_type('torch.FloatTensor')
+
+    def test_per_channel_scale(self):
+        """ fake_tensor_quant performs per channel quantization
+        """
+        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)
+        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
+        # so it has to be reduced by multiple torch.max
+        amax_x_torch = 0.7 * torch.max(
+            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]
+
+        quant_x_np = test_utils.quant_np(x_np, amax_x_np)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, amax_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
+        np.testing.assert_array_less(np.abs(quant_x_torch.cpu().numpy() - quant_x_np), 2)
+        if verbose:
+            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1)
+            print("Mismatches:")
+            print(" Original: ", x_np[mismatches])
+            print(" numpy: ", quant_x_np[mismatches])
+            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])
+
+    def test_backward(self):
+        """ tensor_quant implements straight through estimator on the backward pass
+            Note: this does not work for integer output_dtype
+        """
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x, _ = tensor_quant.tensor_quant(x, x.abs().max(), 7)
+        float_quant_x = quant_x.type(torch.FloatTensor).cuda()
+        x.retain_grad()
+        float_quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(float_quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(float_quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())
+
+    def test_unsigned(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np)
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=False)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+        x_torch = torch.randn(3, 7)
+        with pytest.raises(TypeError, match="Negative values encountered"):
+            tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+
+    def test_overflow_fp16(self):
+        x_torch = torch.randn(1023).cuda().half()
+        with pytest.raises(ValueError, match="scale is too large for FP16"):
+            quant_x_torch, scale = tensor_quant.tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        amax = x.abs().max() / 2
+        x_in_range = (-amax <= x) * (x <= amax)
+        quant_x, _ = tensor_quant.tensor_quant(x, amax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+    def test_full_range(self):
+        """ fake_tensor_quant uses the full integer range when narrow=False
+        """
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = np.max(np.abs(x_np))
+        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=False, narrow_range=False)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+
+class TestFakeTensorQuant():
+
+    def test_simple_run(self):
+        x = Parameter(torch.randn(3, 7).cuda())
+        tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)))
+
+    def test_per_tensor_scale(self):
+        """ fake_tensor_quant matches numpy quantization
+        """
+        x_np = np.random.rand(13).astype('float32')
+        print(x_np)
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    def test_per_channel_scale(self):
+        """ fake_tensor_quant performs per channel quantization
+        """
+        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.9 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
+        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
+        # so it has to be reduced by multiple torch.max
+        amax_x_torch = 0.9 * torch.max(
+            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]
+
+        quant_x_np = test_utils.quant_np(x_np, amax_x_np, fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, amax_x_torch)
+
+        # Pytorch numerics is not the same as numpy, results will be off a little
+        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np, decimal=2)
+        if verbose:
+            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1e-5)
+            print("Mismatches:")
+            print(" Original: ", x_np[mismatches])
+            print(" numpy: ", quant_x_np[mismatches])
+            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])
+
+    def test_backward(self):
+        """ fake_tensor_quant implements straight through estimator on the backward pass
+        """
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x = tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)), 7)
+        x.retain_grad()
+        quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())
+
+    def test_unsigned(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    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 and bf16
+        for dtype in [torch.float16, torch.bfloat16]:
+            x_np = np.random.rand(1023)
+            x_torch = torch.Tensor(x_np).cuda().to(dtype)
+            cuda_ext_out = cuda_ext.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
+            pytorch_out = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
+            test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_cuda_ext_with_axis(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6)
+        x_torch = torch.Tensor(x_np).cuda().to(dtype)
+
+        # 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).to(torch.float32)
+                pytorch_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch.view(1, -1, 1, 1), num_bits, unsigned).to(torch.float32)
+                test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)
+
+    def test_cuda_ext_inplace(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+        cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_equal(x_torch.cpu().numpy(), quant_x_np)
+
+        # Test fp16 and bf16
+        for dtype in [torch.float16, torch.bfloat16]:
+            x_np = np.random.rand(1023)
+            x_torch = torch.Tensor(x_np).cuda().to(dtype)
+            quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+            cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
+            x_torch = x_torch.to(torch.float32)
+            np.testing.assert_array_almost_equal(x_torch.cpu().numpy(), quant_x_np, decimal=2)
+
+    def test_cuda_ext_tiny_amax(self):
+        x_torch = torch.rand(2, 3, 4, device="cuda")
+        amax = torch.tensor([1., 1.e-26, 1.], device="cuda").unsqueeze(-1).unsqueeze(1)
+        quant_x = cuda_ext.fake_tensor_quant_with_axis(x_torch, amax, axis=1)
+        assert quant_x[:, 1, :].sum() == 0
+
+    def test_overflow_fp16(self):
+        x_torch = torch.randn(1023).cuda().half()
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)
+        assert not (torch.isinf(quant_x_torch).any() or torch.isnan(quant_x_torch).any())
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        amax = x.abs().max() / 2
+        x_in_range = (-amax <= x) * (x <= amax)
+        quant_x = tensor_quant.fake_tensor_quant(x, amax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+    def test_full_range(self):
+        """ fake_tensor_quant uses the full integer range when narrow=False
+        """
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = np.max(np.abs(x_np))
+        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=True, narrow_range=False)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    @pytest.mark.parametrize("dtype", ["float32", "float16"])
+    def test_against_legacy(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
+        x_torch = torch.Tensor(x_np).cuda()
+
+        amax_torch = torch.tensor(0.7, device="cuda")
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+    def test_against_legacy_noncontiguous(self):
+        x_np = np.random.rand(3, 4, 5, 6)
+        x_torch = torch.Tensor(x_np).cuda()
+
+        amax_torch = torch.tensor(0.7, device="cuda")
+
+        x_torch_noncontiguous = x_torch[:, 2, :, 3]
+        assert not x_torch_noncontiguous.is_contiguous()
+
+        legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch_noncontiguous, amax_torch)
+        test_out = tensor_quant.fake_tensor_quant(x_torch_noncontiguous, amax_torch)
+        test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("dtype", ["float32", "float16"])
+    def test_against_legacy_with_axis(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
+        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").view(1, -1, 1, 1)
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+
+class TestQuantDescriptor():
+
+    def test_scaled_mode(self):
+        num_bits = np.random.randint(0, 16)
+
+        test_quant_desc = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        assert test_quant_desc.num_bits == num_bits
+        assert test_quant_desc.axis is None
+        assert test_quant_desc.amax is None
+        assert not test_quant_desc.learn_amax
+
+        axis = (0, 1, 3)
+        test_quant_desc = tensor_quant.QuantDescriptor(axis=axis)
+        assert test_quant_desc.num_bits == 8  # default value
+        assert test_quant_desc.axis == axis
+        assert test_quant_desc.amax is None
+
+        amax = 0.7
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, unsigned=True)
+        assert test_quant_desc.axis is None
+        assert test_quant_desc.amax == np.float32(amax)
+        assert test_quant_desc.unsigned
+
+        amax = 0.7
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, learn_amax=True)
+        assert test_quant_desc.amax == np.float32(amax)
+        assert test_quant_desc.learn_amax
+
+        # Test the print string once if verbose is set.
+        if verbose:
+            print(test_quant_desc)
+
+        with pytest.raises(TypeError, match="must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops')
+
+        with pytest.raises(TypeError, match="amax must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops', learn_amax=True)
+
+        with pytest.raises(TypeError, match="axis is ignored and must be None"):
+            tensor_quant.QuantDescriptor(axis=(1, 2), amax=0.7, learn_amax=True)
+
+    def test_amax(self):
+        test_quant_desc = tensor_quant.QuantDescriptor()
+        assert test_quant_desc.amax is None
+
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=1.2)
+        assert isinstance(test_quant_desc.amax, np.ndarray)
+        np.testing.assert_array_equal(test_quant_desc.amax, np.float32(1.2))
+
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=[1.3, 1.4])
+        assert isinstance(test_quant_desc.amax, np.ndarray)
+        np.testing.assert_array_equal(test_quant_desc.amax, np.float32([1.3, 1.4]))
+
+        with pytest.raises(TypeError, match="must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops')
+
+    def test_from_to_dict(self):
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
+                                                    name='a',
+                                                    fake_quant=True,
+                                                    axis=(1, 2),
+                                                    amax=3.1415926536)
+        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
+        if verbose:
+            print(quant_desc_1.dict())
+        assert quant_desc_1 == quant_desc_2
+
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1, unsigned=True)
+        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
+        assert quant_desc_1 == quant_desc_2
+
+    def test_from_to_yaml(self):
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
+                                                    name='a',
+                                                    fake_quant=True,
+                                                    axis=(1, 2),
+                                                    amax=3.1415926536)
+        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
+        if verbose:
+            print(quant_desc_1.to_yaml())
+        assert quant_desc_1 == quant_desc_2
+
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1)
+        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
+        assert quant_desc_1 == quant_desc_2
+
+
+class TestFakeAffineTensorQuant():
+
+    def test_simple_run(self, verbose):
+        x = np.array([-1., -13., -101., -128., 0., 2., 5., 13., 93., 111., 127.], dtype=np.float32)
+        torch_x = torch.tensor(x).cuda()
+        quant_x = tensor_quant.fake_affine_tensor_quant(torch_x, torch.min(torch_x), torch.max(torch_x))
+
+        if verbose:
+            print(quant_x)
+
+        np.testing.assert_array_almost_equal(quant_x.cpu().numpy(), x)
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        xmin = x.min() / 2
+        xmax = x.max() / 2
+        x_in_range = (xmin <= x) * (x <= xmax)
+        quant_x = tensor_quant.fake_affine_tensor_quant(x, xmin, xmax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+
+class TestScaledE4M3():
+
+    x = [[-2.0000, -1.8000, -1.6000, -1.4000, -1.2000], [-1.0000, -0.8000, -0.6000, -0.4000, -0.2000],
+         [-0.0000, 0.2000, 0.4000, 0.6000, 0.8000], [1.0000, 1.2000, 1.4000, 1.6000, 1.8000]]
+
+    xq_unscaled = [[-2.0000, -1.7500, -1.6250, -1.3750, -1.2500], [-1.0000, -0.8125, -0.6250, -0.4062, -0.2031],
+                   [0.0000, 0.2031, 0.4062, 0.6250, 0.8125], [1.0000, 1.2500, 1.3750, 1.6250, 1.7500]]
+
+    xq_scaled = [[-2.0000, -1.8571, -1.5714, -1.4286, -1.1429], [-1.0000, -0.7857, -0.5714, -0.3929, -0.1964],
+                 [0.0000, 0.1964, 0.3929, 0.5714, 0.7857], [1.0000, 1.1429, 1.4286, 1.5714, 1.8571]]
+
+    def test_e4m3_no_scale(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda")
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda")
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_e4m3_no_cpu(self):
+        x = torch.tensor(TestScaledE4M3.x)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled)
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_with_amax(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda").unsqueeze(-1)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_scaled, device="cuda").unsqueeze(-1)
+
+        amax = quant_utils.reduce_amax(x, axis=None, keepdims=True)
+
+        e4m3_x = tensor_quant.scaled_e4m3(x, amax)
+
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_e4m3_incontiguous(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda").transpose(1, 0)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda").transpose(1, 0)
+        assert not x.is_contiguous()
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_backward(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x = tensor_quant.scaled_e4m3(x, None)
+        x.retain_grad()
+        quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())

tools/pytorch-quantization/tests/tensor_quantizer_test.py

@@ -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)

tools/pytorch-quantization/tests/test_onnx_export.py

@@ -0,0 +1,111 @@
+#
+# 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 for ONNX export."""
+import io
+import onnxruntime
+import pytest
+import torch
+
+# ORT output correctness tests sometimes fails due to random seed.
+# It needs to be investigated closer
+torch.manual_seed(0)
+
+import tests.utils as test_utils
+import torch.nn as nn
+import pytorch_quantization
+from pytorch_quantization.nn import QuantLinear
+from pytorch_quantization.tensor_quant import QuantDescriptor
+
+
+class MyModel(nn.Module):
+    """Test model for ONNX export."""
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.net = nn.Sequential(
+            QuantLinear(16, 32, **kwargs),
+            nn.ReLU(),
+            QuantLinear(32, 64, **kwargs),
+            nn.ReLU(),
+            QuantLinear(64, 16, **kwargs),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+@pytest.mark.parametrize("num_bits, per_channel_quantization, constant_folding, dtype",
+                         [(8, True, True, torch.float32), (8, False, True, torch.float32),
+                          (8, True, False, torch.float32), (8, False, False, torch.float32),
+                          (8, False, False, torch.float16), (8, False, False, torch.bfloat16),
+                          ((4, 3), False, True, torch.float32), ((4, 3), False, False, torch.float32),
+                          ((4, 3), False, False, torch.float16), ((4, 3), False, False, torch.bfloat16)])
+def test_onnx_export(num_bits, per_channel_quantization, constant_folding, dtype, onnx_file_path=None):
+    quant_desc_input = QuantDescriptor(num_bits=num_bits, axis=None)
+    quant_desc_weight = QuantDescriptor(num_bits=num_bits, axis=0 if per_channel_quantization else None)
+
+    model = MyModel(quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight).cuda()
+    model.eval()
+
+    OPSET = 17
+    dummy_input = torch.randn(16, 16).cuda()
+    input_names = ["input"]
+    output_names = ["output"]
+
+    model = model.to(dtype)
+    dummy_input = dummy_input.to(dtype)
+
+    # Calibrate model
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.enable_calib()
+            module.disable_quant()
+    _ = model(dummy_input)
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.disable_calib()
+            module.load_calib_amax()
+            module.enable_quant()
+
+    f = io.BytesIO() if onnx_file_path is None else None
+
+    with pytorch_quantization.enable_onnx_export():
+        torch.onnx.export(
+            model,
+            dummy_input,
+            f=f if onnx_file_path is None else onnx_file_path,
+            opset_version=OPSET,
+            input_names=input_names,
+            output_names=output_names,
+            do_constant_folding=constant_folding,
+        )
+
+    # TODO: ort output correctness check for fp8
+    # ONNXRuntime does not seem to be supporting bf16 gemms
+    if num_bits == 8 and dtype != torch.bfloat16:
+        if f is not None:
+            f.seek(0)
+        ort_session = onnxruntime.InferenceSession(f.read() if onnx_file_path is None else onnx_file_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        ort_result = ort_session.run([], {"input": dummy_input.cpu().numpy()})
+        ort_result = torch.tensor(ort_result[0]).cuda()
+        torch_result = model(dummy_input)
+        test_utils.compare(ort_result, torch_result, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    test_onnx_export(8, False, False, torch.float16, "/tmp/test_fp16.onnx")
+    test_onnx_export(8, False, False, torch.bfloat16, "/tmp/test_bf16.onnx")

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